Leetcode 1989 Maximum Number of People That Can Be Caught in Tag

Problem statement

https://leetcode.com/problems/maximum-number-of-people-that-can-be-caught-in-tag/

Solution

The idea is to at each moment of time look at the most left 1 and at the most left 0. If we can combine them, we do it, if not, we move the one which is behind. We can create two queues and extract element from the left (or use stack where we put reversed data) Why it is working? Because for the leftest 1 it is always better to choose the leftest 0.

Complexity

It is O(n) both for time and space.

Code

class Solution:
    def catchMaximumAmountofPeople(self, team, dist):
        zeroes = deque([i for i, num in enumerate(team) if num == 0])
        ones = deque([i for i, num in enumerate(team) if num == 1])
        ans = 0
        while zeroes and ones:
            if abs(zeroes[0] - ones[0]) <= dist:
                ans += 1
                zeroes.popleft()
                ones.popleft()
            elif zeroes[0] < ones[0]:
                zeroes.popleft()
            else:
                ones.popleft()
        
        return ans