All problems
RAG Retrieval Top-k
Problem
Return the top-k passage ids for a query. passages is a list of [id, vector] pairs.
Score each passage by the dot product with query, then sort by descending score with ties broken alphabetically by id. Return an empty list when k <= 0.
Hint
Sort by the tuple (negative score, id) to get the tie-breaking for free.
Test cases
case 1
in: [[1,0],[["a",[1,0]],["b",[0,1]],["c",[2,0]]],2]
expected: ["c","a"]
case 2
in: [[1,1],[["x",[1,1]],["y",[0,1]]],1]
expected: ["x"]
case 3
in: [[1,0],[["b",[0,1]],["a",[0,1]]],2]
expected: ["a","b"]
case 4
in: [[1,0],[],3]
expected: []
case 5
in: [[1,0],[["a",[1,0]]],0]
expected: []
Starter codePython
def rag_topk(query, passages, k):
# Your code here
pass5 test cases · run your code against them in the editor