2528-最大化城市的最小供电站数目

给你一个下标从 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

前置知识 1：二分

见 【基础算法精讲 04】 。

前置知识 2&3：前缀和、差分数组

见 视频讲解 的第四题。

提示 1

看到「最大化最小值」或者「最小化最大值」就要想到二分答案，这是一个固定的套路。

为什么？一般来说，二分的值越大，越能/不能满足要求；二分的值越小，越不能/能满足要求，有单调性，可以二分。

类似的题目在先前的周赛中出现过多次，例如：

• 2439. 最小化数组中的最大值
• 2513. 最小化两个数组中的最大值
• 2517. 礼盒的最大甜蜜度

提示 2

二分答案 minPower，从左到右遍历 stations，如果 stations}[i] 电量不足 minPower，那么需要建供电站来补足。

在哪建供电站最好呢？

提示 3

由于 i 左侧的不需要补足，所以贪心地在 \min(i+r,n-1) 处建是最合适的，恰好让 i 在覆盖范围的边界上。

提示 4

设需要建 m 个供电站，那么需要把 [i,\min(i+2r,n-1)] 范围内的电量都增加 m。

方法很多，用差分数组来更新是最简单的。

最后判断修建的供电站是否超过 k，如果超过说明 minPower 偏大，否则说明偏小。

注：其实前缀和也不需要，可以改为长为 2r+1 的滑动窗口，但这样写有点麻烦，感兴趣的读者可以实现下。

[sol1-Python3]

class Solution:
    def maxPower(self, stations: List[int], r: int, k: int) -> int:
        n = len(stations)
        sum = list(accumulate(stations, initial=0))  # 前缀和
        for i in range(n):
            stations[i] = sum[min(i + r + 1, n)] - sum[max(i - r, 0)]  # 电量

        def check(min_power: int) -> bool:
            diff = [0] * n  # 差分数组
            sum_d = need = 0
            for i, power in enumerate(stations):
                sum_d += diff[i]  # 累加差分值
                m = min_power - power - sum_d
                if m > 0:  # 需要 m 个供电站
                    need += m
                    if need > k: return False  # 提前退出这样快一些
                    sum_d += m  # 差分更新
                    if i + r * 2 + 1 < n: diff[i + r * 2 + 1] -= m  # 差分更新
            return True

        left = min(stations)
        right = left + k + 1  # 开区间写法
        while left + 1 < right:
            mid = (left + right) // 2
            if check(mid): left = mid
            else: right = mid
        return left

[sol1-Java]

class Solution {
    public long maxPower(int[] stations, int r, int k) {
        int n = stations.length;
        long[] sum = new long[n + 1]; // 前缀和
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + stations[i];
        long mn = Long.MAX_VALUE;
        long[] power = new long[n]; // 电量
        for (int i = 0; i < n; ++i) {
            power[i] = sum[Math.min(i + r + 1, n)] - sum[Math.max(i - r, 0)];
            mn = Math.min(mn, power[i]);
        }

        long left = mn, right = mn + k + 1; // 开区间写法
        while (left + 1 < right) {
            long mid = left + (right - left) / 2;
            if (check(mid, power, n, r, k)) left = mid;
            else right = mid;
        }
        return left;
    }

    private boolean check(long minPower, long[] power, int n, int r, int k) {
        long[] diff = new long[n + 1]; // 差分数组
        long sumD = 0, need = 0;
        for (int i = 0; i < n; ++i) {
            sumD += diff[i]; // 累加差分值
            long m = minPower - power[i] - sumD;
            if (m > 0) { // 需要 m 个供电站
                need += m;
                if (need > k) return false; // 提前退出这样快一些
                sumD += m; // 差分更新
                if (i + r * 2 + 1 < n) diff[i + r * 2 + 1] -= m; // 差分更新
            }
        }
        return true;
    }
}

[sol1-C++]

class Solution {
public:
    long long maxPower(vector<int> &stations, int r, int k) {
        int n = stations.size();
        long sum[n + 1], power[n], diff[n];
        sum[0] = 0;
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + stations[i]; // 前缀和
        for (int i = 0; i < n; ++i)
            power[i] = sum[min(i + r + 1, n)] - sum[max(i - r, 0)]; // 电量

        auto check = [&](long min_power) -> bool {
            memset(diff, 0, sizeof(diff)); // 重置差分数组
            long sum_d = 0, need = 0;
            for (int i = 0; i < n; ++i) {
                sum_d += diff[i]; // 累加差分值
                long m = min_power - power[i] - sum_d;
                if (m > 0) { // 需要 m 个供电站
                    need += m;
                    if (need > k) return false; // 提前退出这样快一些
                    sum_d += m; // 差分更新
                    if (i + r * 2 + 1 < n) diff[i + r * 2 + 1] -= m; // 差分更新
                }
            }
            return true;
        };

        long left = *min_element(power, power + n), right = left + k + 1; // 开区间写法
        while (left + 1 < right) {
            long mid = left + (right - left) / 2;
            check(mid) ? left = mid : right = mid;
        }
        return left;
    }
};

[sol1-Go]

func maxPower(stations []int, r int, k int) int64 {
	n := len(stations)
	sum := make([]int, n+1) // 前缀和
	for i, x := range stations {
		sum[i+1] = sum[i] + x
	}
	mn := math.MaxInt
	for i := range stations {
		stations[i] = sum[min(i+r+1, n)] - sum[max(i-r, 0)] // 电量
		mn = min(mn, stations[i])
	}
	return int64(mn + sort.Search(k, func(minPower int) bool {
		minPower += mn + 1 // 改为二分最小的不满足要求的值，这样 sort.Search 返回的就是最大的满足要求的值
		diff := make([]int, n) // 差分数组
		sumD, need := 0, 0
		for i, power := range stations {
			sumD += diff[i] // 累加差分值
			m := minPower - power - sumD
			if m > 0 { // 需要 m 个供电站
				need += m
				if need > k { // 提前退出这样快一些
					return true // 不满足要求
				}
				sumD += m // 差分更新
				if i+r*2+1 < n {
					diff[i+r*2+1] -= m // 差分更新
				}
			}
		}
		return false
	}))
}

func min(a, b int) int { if b < a { return b }; return a }
func max(a, b int) int { if b > a { return b }; return a }

复杂度分析

• 时间复杂度：O(n\log k)，其中 n 为 stations 的长度。二分需要循环 O(\log k) 次。
• 空间复杂度：O(n)。