最大化城市的最小电量

标签: 贪心 队列 数组 二分查找 前缀和 滑动窗口

难度: Hard

给你一个下标从 0 开始长度为 n 的整数数组 stations ,其中 stations[i] 表示第 i 座城市的供电站数目。

每个供电站可以在一定 范围 内给所有城市提供电力。换句话说,如果给定的范围是 r ,在城市 i 处的供电站可以给所有满足 |i - j| <= r 且 0 <= i, j <= n - 1 的城市 j 供电。

  • |x| 表示 x 的 绝对值 。比方说,|7 - 5| = 2 ,|3 - 10| = 7 。

一座城市的 电量 是所有能给它供电的供电站数目。

政府批准了可以额外建造 k 座供电站,你需要决定这些供电站分别应该建在哪里,这些供电站与已经存在的供电站有相同的供电范围。

给你两个整数 r 和 k ,如果以最优策略建造额外的发电站,返回所有城市中,最小电量的最大值是多少。

k 座供电站可以建在多个城市。

示例 1:

输入:stations = [1,2,4,5,0], r = 1, k = 2
输出:5
解释:
最优方案之一是把 2 座供电站都建在城市 1 。
每座城市的供电站数目分别为 [1,4,4,5,0] 。
- 城市 0 的供电站数目为 1 + 4 = 5 。
- 城市 1 的供电站数目为 1 + 4 + 4 = 9 。
- 城市 2 的供电站数目为 4 + 4 + 5 = 13 。
- 城市 3 的供电站数目为 5 + 4 = 9 。
- 城市 4 的供电站数目为 5 + 0 = 5 。
供电站数目最少是 5 。
无法得到更优解,所以我们返回 5 。

示例 2:

输入:stations = [4,4,4,4], r = 0, k = 3
输出:4
解释:
无论如何安排,总有一座城市的供电站数目是 4 ,所以最优解是 4 。

提示:

  • n == stations.length
  • 1 <= n <= 105
  • 0 <= stations[i] <= 105
  • 0 <= r <= n - 1
  • 0 <= k <= 109

Submission

运行时间: 864 ms

内存: 29.1 MB

class Solution:
    def maxPower(self, stations: List[int], r: int, k: int) -> int:
        # 二分答案 + check
        n = len(stations)
        pre = [0] * n
        left = 0
        sm = 0
        if r >= n - 1:
            return k + sum(stations)
        else:
            for i in range(r):
                sm += stations[i]
            
            for i in range(r, n):
                if i - left >  2 * r:
                    sm -= stations[left]
                    left += 1
                sm += stations[i]
                pre[i - r] = sm
            for i in range(n, n + r):
                if i - left >  2 * r:
                    sm -= stations[left]
                    left += 1
                pre[i - r] = sm
        

        print(pre)
        def check(x :int) -> bool:
            # 查分数组
            diff = [0] * n
            cur = 0
            cnt = 0
            for i, el in enumerate(pre):
                cur += diff[i]
                if el + cur < x:
                    addi = (x - cur - el)
                    cur += addi
                    cnt += addi
                    if cnt > k:
                        return False
                    if 2 * r + 1 + i < n:
                        diff[2 * r + 1 + i] -= addi
            return True



        
        l = 0
        rr = k + sum(stations)
        while l < rr:
            mid = l + (rr - l)//2
            if check(mid):
                l = mid + 1
            else:
                rr = mid -1
        return l if check(l) else l - 1

Explain

题解采用了二分查找和差分数组的组合技术来解决问题。首先,根据电站的覆盖范围 r,初始化一个前缀和数组 pre,用于计算每个城市的初始供电情况。如果 r 覆盖整个数组,则简单地将 k 加到所有电站数之和上。否则,使用滑动窗口计算每个城市由覆盖范围内的电站提供的电力总和。接下来,使用二分搜索尝试找到可能的最小电量的最大值。对于二分搜索中的每个中间值 mid,使用差分数组来模拟在不同位置增加 k 个电站后的电力分布,以检查是否所有城市的电量都可以达到 mid。最终通过调整二分搜索的上下界,找到最大的可行电量。

时间复杂度: O(n log(maxPower))

空间复杂度: O(n)

class Solution:
    def maxPower(self, stations: List[int], r: int, k: int) -> int:
        # 初始化
        n = len(stations)
        pre = [0] * n # 存储覆盖每个城市的电站总数
        left = 0 # 滑动窗口的左边界
        sm = 0 # 窗口内电站的累计数
        if r >= n - 1:
            return k + sum(stations) # 如果覆盖范围超过城市总数,直接返回
        else:
            # 构建前缀和数组
            for i in range(r):
                sm += stations[i]
            for i in range(r, n):
                if i - left >  2 * r:
                    sm -= stations[left]
                    left += 1
                sm += stations[i]
                pre[i - r] = sm
            for i in range(n, n + r):
                if i - left >  2 * r:
                    sm -= stations[left]
                    left += 1
                pre[i - r] = sm
        
        # 二分查找和检查函数
        def check(x: int) -> bool:
            diff = [0] * n
            cur = 0 # 当前增加的电站数
            cnt = 0 # 总共增加的电站数
            for i, el in enumerate(pre):
                cur += diff[i]
                if el + cur < x: # 如果当前电量低于mid
                    addi = (x - cur - el) # 需要增加的电站数
                    cur += addi
                    cnt += addi
                    if cnt > k:
                        return False # 如果超过可用电站数,返回 False
                    if 2 * r + 1 + i < n:
                        diff[2 * r + 1 + i] -= addi # 更新差分数组
            return True
        l = 0
        rr = k + sum(stations) # 最大可能电量
        while l < rr:
            mid = l + (rr - l) // 2
            if check(mid):
                l = mid + 1 # 尝试更大的 mid
            else:
                rr = mid -1 # 降低 mid
        return l if check(l) else l - 1 # 最后检查 l 是否有效

Explore

二分搜索被选择用于解决此问题,因为我们需要在可能的电量值范围内找到一个最大值,使得所有城市的电量都不低于这个值。这是一个典型的“决策问题”,我们需要判断一个电量值是否能满足条件。通过不断地调整电量值范围来逼近最大可能的电量值,二分搜索能有效地在有序的数据范围中快速缩小搜索区间,从而找到最优解。

在使用差分数组进行模拟时,算法会逐一检查每个城市当前的电量是否低于目标中值(mid)。如果某个城市的电量不足,算法计算出必须增加的电站数,并累加到总增加的电站数(cnt)中。每当增加电站时,都会检查累加后的总电站数是否超过了k。如果超过k,则函数返回False,表示当前的中值不可行。这样,通过在增加电站数时持续监控总数,我们确保不会超出k的限制。同时,差分数组的更新确保了增加的电站影响可以正确地延伸到适当的范围内,而不会错误地扩展到数组外。

差分数组是一种用于表示数值序列中相邻元素差值的技术,可以高效地处理区间增减值的问题。在本题中,当发现某个城市的电量低于目标值时,我们不仅需要在当前位置增加必要的电站,还需要确保这种增加影响可以覆盖到电站的最大范围(即2r)。通过在当前位置增加电站数,并在超出电站范围的下一个位置减去相同的电站数,差分数组能够有效地模拟每个城市实际受到的电站影响。随后,通过累加差分数组来得到每个城市实际的电站数,这种方法使得多次更新操作的复杂度仍然保持在线性级别。