class KthLargest: def __init__(self, k, nums): self.k = k self.pq = PriorityQueue() for i in nums: self.pq.put(i) def add(self, n): if self.pq.qsize() < self.k: self.pq.put(n) return self.pq.get_nowait()
def printKclosest(arr, n, x, k): # Make a max heap of difference with # first k elements. pq = PriorityQueue() for i in range(k): pq.put((-abs(arr[i] - x), i)) # Now process remaining elements for i in range(k, n): diff = abs(arr[i] - x) p, pi = pq.get() curr = -p # If difference with current # element is more than root, # then put it back. if diff > curr: pq.put((-curr, pi)) continue else: # Else remove root and insert pq.put((-diff, i)) # Print contents of heap. curr = 0 while (not pq.empty()): p, q = pq.get() curr += abs(val - arr[q]) last_ans[val] = curr
def beam_decode(target_tensor, decoder_hiddens, encoder_outputs=None): ''' :param target_tensor: target indexes tensor of shape [B, T] where B is the batch size and T is the maximum length of the output sentence :param decoder_hidden: input tensor of shape [1, B, H] for start of the decoding :param encoder_outputs: if you are using attention mechanism you can pass encoder outputs, [T, B, H] where T is the maximum length of input sentence :return: decoded_batch ''' beam_width = 10 topk = 1 # how many sentence do you want to generate decoded_batch = [] # decoding goes sentence by sentence for idx in range(target_tensor.size(0)): if isinstance(decoder_hiddens, tuple): # LSTM case decoder_hidden = (decoder_hiddens[0][:, idx, :].unsqueeze(0), decoder_hiddens[1][:, idx, :].unsqueeze(0)) else: decoder_hidden = decoder_hiddens[:, idx, :].unsqueeze(0) encoder_output = encoder_outputs[:, idx, :].unsqueeze(1) # Start with the start of the sentence token decoder_input = torch.LongTensor([[SOS_token]], device=device) # Number of sentence to generate endnodes = [] number_required = min((topk + 1), topk - len(endnodes)) # starting node - hidden vector, previous node, word id, logp, length node = BeamSearchNode(decoder_hidden, None, decoder_input, 0, 1) nodes = PriorityQueue() # start the queue nodes.put((-node.eval(), node)) qsize = 1 # start beam search while True: # give up when decoding takes too long if qsize > 2000: break # fetch the best node score, n = nodes.get() decoder_input = n.wordid decoder_hidden = n.h if n.wordid.item() == EOS_token and n.prevNode != None: endnodes.append((score, n)) # if we reached maximum # of sentences required if len(endnodes) >= number_required: break else: continue # decode for one step using decoder decoder_output, decoder_hidden = decoder(decoder_input, decoder_hidden, encoder_output) # PUT HERE REAL BEAM SEARCH OF TOP log_prob, indexes = torch.topk(decoder_output, beam_width) nextnodes = [] for new_k in range(beam_width): decoded_t = indexes[0][new_k].view(1, -1) log_p = log_prob[0][new_k].item() node = BeamSearchNode(decoder_hidden, n, decoded_t, n.logp + log_p, n.leng + 1) score = -node.eval() nextnodes.append((score, node)) # put them into queue for i in range(len(nextnodes)): score, nn = nextnodes[i] nodes.put((score, nn)) # increase qsize qsize += len(nextnodes) - 1 # choose nbest paths, back trace them if len(endnodes) == 0: endnodes = [nodes.get() for _ in range(topk)] utterances = [] for score, n in sorted(endnodes, key=operator.itemgetter(0)): utterance = [] utterance.append(n.wordid) # back trace while n.prevNode != None: n = n.prevNode utterance.append(n.wordid) utterance = utterance[::-1] utterances.append(utterance) decoded_batch.append(utterances) return decoded_batch