Problem Statement in English

You’re given a 0-indexed array stations of length n, where stations[i] represents the number of power stations in the ith city.

Each power station in a city provides power to all cities within a range r. In other words, a power station at city i provides power to all cities j such that |i - j| <= r.

You are also given an integer k, which is the maximum number of additional power stations that can be built. You can build the additional power stations in any city, including cities that already have power stations.

The power of a city is defined as the total number of power stations that provide power to it. The minimum power of a city is the smallest power among all the cities.

Return the maximum possible minimum power of a city after building at most k additional power stations.

Approach

I highly recommend watching NeetCode’s explanation of this problem. It is very well done.

There are 2 main parts that we need to solve this problem.

The first is to efficiently represent and compute the power of each city - whilst factoring in the range r of each power station. If you’ve done problems on the sweep line algorithm, you might figure that part out easily enough.

The next part, which is also common, although not super common, is the what I like to call the “guess binary search”. You simply take a guess of what the answer might be, and check if it’s feasible to achieve that answer. If it is, you try for a higher guess. If it isn’t, you try for a lower guess.

To do this, you need to know the answer space. In our case, the minimum power of any city (and hence, the left boundary of our answer space) is at least the minimum power in the original stations array. The maximum power of any city is at most the sum of all powers in stations plus k (imagine that r encapsulates all cities in our array - this means that increasing the power of one city extends that power to all other cities). Read it again if it doesn’t make sense. It is vital to our solution and even being able to do this whole guess based binary search.

The other part of the problem is the difference array. We just increase the power at the start of the range, and decrease it just after the end of the range. Read that portion of the code if that didn’t make sense.

One quick thing to note here is being careful with the bounds when updating the difference array. We need to ensure that we don’t go out of bounds of the array. Not just that, we need to ensure that we don’t decrease the power at the last city when the power actually extends beyond the last city. That’s why we extend the length of the difference array by 1.

It’s not so bad once you visualize it in your head a couple times.

Solution in Python


class Solution:
    def maxPower(self, stations: List[int], r: int, k: int) -> int:
        N = len(stations)

        diffArray = [0] * (N + 1)

        for i, v in enumerate(stations):
            diffArray[max(0, i - r)] += v
            diffArray[min(N, i + r + 1)] -= v
        
        lowest = min(stations)
        highest = sum(stations) + k

        def can_achieve(guess):
            copied = diffArray.copy()
            ops = k
            current_power = 0

            for i in range(N):
                current_power += copied[i]

                if current_power < guess:
                    opsNeeded = guess - current_power

                    if opsNeeded > ops:
                        return False

                    current_power += opsNeeded
                    copied[min(N, i + 2 * r + 1)] -= opsNeeded
                    ops -= opsNeeded

            return True

        left, right = lowest, highest
        res = left

        while left <= right:
            mid = (left + right) // 2

            if can_achieve(mid):
                res = mid
                left = mid + 1
            else:
                right = mid - 1

        return res

Complexity

  • Time: $O(N \log(S))$
    Here, $N$ is the number of stations and $S$ is the sum of all station powers plus $k$.

  • Space: $O(N)$
    For the difference array used in the solution.

Mistakes I Made

I watched NeetCode’s solution, and wow. Full credit to him. I would not have attempted this problem had it not been for the fact that he’s making videos again.

And we are done.