Leetcode 2012 Solution

class Solution:
    def assignTasks(self, servers: List[int], tasks: List[int]) -> List[int]:
        pending_servers = []
        active_servers = []

        for i, weight in enumerate(servers):
            heapq.heappush(active_servers, (weight, i))

        t = 0
        i = 0
        ans = [0] * len(tasks)
        while i < len(tasks):
            while pending_servers and pending_servers[0][0] <= t:
                _, weight, server_id = heapq.heappop(pending_servers)
                heapq.heappush(active_servers, (weight, server_id))

            while active_servers and i <= t and i < len(tasks):
                weight, server_id = heapq.heappop(active_servers)

                ans[i] = server_id
                heapq.heappush(pending_servers, (t + tasks[i], weight, server_id))

                i += 1

            if not active_servers:
                t = pending_servers[0][0]
            else:
                t += 1

        return ans