Merge K Sorted Lists

08 Nov 2018 in Programming on Heap, Interview, Python

Merge k sorted linked lists and return it as one sorted list.

Example :

1 -> 10 -> 20
4 -> 11 -> 13
3 -> 8 -> 9

will result in

1 -> 3 -> 4 -> 8 -> 9 -> 10 -> 11 -> 13 -> 20

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None
import heapq
class Solution:
    # @param A : list of linked list
    # @return the head node in the linked list
        def mergeKLists(self, lists):
        h = [(l.val, idx) for idx, l in enumerate(lists) if l]
        heapq.heapify(h)
        head = cur = ListNode(None)
        while h:
            val, idx = heapq.heappop(h)
            cur.next = ListNode(val)
            cur = cur.next
            node = lists[idx] = lists[idx].next
            if node:
                heapq.heappush(h, (node.val, idx))
        return head.next

Magician and Chocolates

Given N bags, each bag contains Ai chocolates. There is a kid and a magician. In one unit of time, kid chooses a random bag i, eats Ai chocolates, then the magician fills the ith bag with floor(Ai/2) chocolates.

Given Ai for 1 <= i <= N, find the maximum number of chocolates kid can eat in K units of time.

For example,

K = 3
N = 2
A = 6 5

Return: 14

At t = 1 kid eats 6 chocolates from bag 0, and the bag gets filled by 3 chocolates At t = 2 kid eats 5 chocolates from bag 1, and the bag gets filled by 2 chocolates At t = 3 kid eats 3 chocolates from bag 0, and the bag gets filled by 1 chocolate so, total number of chocolates eaten: 6 + 5 + 3 = 14

Note: Return your answer modulo 10**9+7

import heapq
import math
class Solution:
    # @param A : integer
    # @param B : list of integers
    # @return an integer
    def nchoc(self, A, B):
        h = [-num for num in B]
        heapq.heapify(h)
        sum = 0
        for _ in range(A):
            val = heapq.heappop(h)
            val = -1 * val
            sum += val
            # sum = sum
            new_val = math.floor(val/2)
            heapq.heappush(h, -1 * new_val)
        return int(sum % (10**9+7))