Problem statement

https://leetcode.com/problems/maximum-number-of-tasks-you-can-assign/

Solution

The idea of this problem is the following: let us ask question: check(k): can we finish k tasks or not. Here is a couple of insights.

  1. If we want to finish k tasks, we better choose the smallest k of them.
  2. If we want to finish k tasks, we need k workers, and we want to choost the k strongest of them.
  3. So, we choose tasks and workers and now the question is can we allocate them. We start from the biggest task and start to allocate it to worker. First we try to give it to worker without pill and if we are OK, we allocate it to the weakest worker. If we can not allocate it to worker without pill, we allocate it to weakest worker with pill. If we can not do both of these steps, we return False: it means, that we were not able to allocate all tasks.

Good question though, why this greedy strategy will work? At the moment I do not have strict proof, I will add it a bit later.

Complexity

It is O(n * log^2 n) for time and O(n) for space.

Code

from sortedcontainers import SortedList

class Solution:
    def maxTaskAssign(self, tasks, workers, pills, strength):
        tasks = sorted(tasks)
        workers = sorted(workers)

        def check(k):
            W = SortedList(workers[-k:])
            tries = pills

            for elem in tasks[:k][::-1]:
                place = W.bisect_left(elem)
                if place < len(W):
                    W.pop(place)
                elif tries > 0:
                    place2 = W.bisect_left(elem - strength)
                    if place2 < len(W):
                        W.pop(place2)
                        tries -= 1
                else:
                    return False

            return len(W) == 0

        beg, end = 0, min(len(workers), len(tasks)) + 1
        while beg + 1 < end:
            mid = (beg + end)//2
            if check(mid):
                beg = mid
            else:
                end = mid

        return beg