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])