class KBest:
def __init__(self,k,x,y):
"""
x and y is a list of tuples or lists -> (score,obj1,obj2)
"""
self.x = x
self.y = y
self.C = PriorityQueue()
self.p = []
self.k = k
self.tmp = []
def _insert_queue(self,i,j):
self.C.put((self.x[i][0]+self.y[j][0],[i,j]))
self.tmp.append([i,j])
def mult(self,k=None):
if k != None:
self.k = k
self._insert_queue(0,0)
while len(self.p) < self.k and not self.C.empty():
self._append_next()
return self.p
def _append_next(self):
s,v = self.C.get()
self.tmp.remove(v)
self.p.append((s,v))
for i in range(2):
vc = copy.deepcopy(v)
vc[i] += 1
if vc[0] < len(self.x) and vc[1] < len(self.y) and vc not in self.tmp:
self._insert_queue(vc[0],vc[1])
