Question

2772. 使数组中的所有元素都等于零

English Version

题目描述

给你一个下标从 0 开始的整数数组 nums 和一个正整数 k 。

你可以对数组执行下述操作 任意次 ：

• 从数组中选出长度为 k 的 任一 子数组，并将子数组中每个元素都 减去 1 。

如果你可以使数组中的所有元素都等于 0 ，返回  true_ _；否则，返回_ _false_ _。

子数组 是数组中的一个非空连续元素序列。

 

示例 1：

输入：nums = [2,2,3,1,1,0], k = 3
输出：true
解释：可以执行下述操作：
- 选出子数组 [2,2,3] ，执行操作后，数组变为 nums = [_1_,_1_,_2_,1,1,0] 。
- 选出子数组 [2,1,1] ，执行操作后，数组变为 nums = [1,1,_1_,_0_,_0_,0] 。
- 选出子数组 [1,1,1] ，执行操作后，数组变为 nums = [_0_,_0_,_0_,0,0,0] 。

示例 2：

输入：nums = [1,3,1,1], k = 2
输出：false
解释：无法使数组中的所有元素等于 0 。

 

提示：

• 1 <= k <= nums.length <= 105
• 0 <= nums[i] <= 106

解法

方法一：差分数组 + 前缀和

我们先考虑 $nums$ 的第一个元素 $nums[0]$：

• 如果 $nums[0] = 0$，那么我们可以不用操作；
• 如果 $nums[0] \gt 0$，那么我们需要对 $nums[0..k-1]$ 操作 $nums[0]$ 次，使得 $nums[0..k-1]$ 中的元素都减去 $nums[0]$，这样 $nums[0]$ 就变成了 $0$。

对一段连续的元素同时进行加减操作，我们可以使用差分数组来维护这些操作，我们用 $d[i]$ 表示差分数组，对差分数组求前缀和，就可以得到每个位置的数值的变化量。

因此，我们遍历 $nums$，对于每个元素 $nums[i]$，当前位置的变化量 $s = \sum_{j=0}^{i} d[j]$，我们将 $nums[i]$ 加上 $s$，就得到了当前 $nums[i]$ 的实际值。

• 如果 $nums[i] = 0$，那么无须进行操作，直接跳过。
• 如果 $nums[i]=0$，或者 $i + k \gt n$，说明经过前面的操作，$nums[i]$ 已经变成了负数，或者 $nums[i..i+k-1]$ 越界，那么无法使得 $nums$ 中的所有元素都等于 $0$，返回 false。否则，我们需要将 $[i..i+k-1]$ 这段区间的所有元素都减去 $nums[i]$，因此我们将 $s$ 减去 $nums[i]$，并将 $d[i+k]$ 加上 $nums[i]$。
• 继续遍历下一个元素。

遍历结束，说明可以使得 $nums$ 中的所有元素都等于 $0$，返回 true。

时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是数组 $nums$ 的长度。

class Solution:
  def checkArray(self, nums: List[int], k: int) -> bool:
    n = len(nums)
    d = [0] * (n + 1)
    s = 0
    for i, x in enumerate(nums):
      s += d[i]
      x += s
      if x == 0:
        continue
      if x < 0 or i + k > n:
        return False
      s -= x
      d[i + k] += x
    return True

class Solution {
  public boolean checkArray(int[] nums, int k) {
    int n = nums.length;
    int[] d = new int[n + 1];
    int s = 0;
    for (int i = 0; i < n; ++i) {
      s += d[i];
      nums[i] += s;
      if (nums[i] == 0) {
        continue;
      }
      if (nums[i] < 0 || i + k > n) {
        return false;
      }
      s -= nums[i];
      d[i + k] += nums[i];
    }
    return true;
  }
}

class Solution {
public:
  bool checkArray(vector<int>& nums, int k) {
    int n = nums.size();
    vector<int> d(n + 1);
    int s = 0;
    for (int i = 0; i < n; ++i) {
      s += d[i];
      nums[i] += s;
      if (nums[i] == 0) {
        continue;
      }
      if (nums[i] < 0 || i + k > n) {
        return false;
      }
      s -= nums[i];
      d[i + k] += nums[i];
    }
    return true;
  }
};

func checkArray(nums []int, k int) bool {
  n := len(nums)
  d := make([]int, n+1)
  s := 0
  for i, x := range nums {
    s += d[i]
    x += s
    if x == 0 {
      continue
    }
    if x < 0 || i+k > n {
      return false
    }
    s -= x
    d[i+k] += x
  }
  return true
}

function checkArray(nums: number[], k: number): boolean {
  const n = nums.length;
  const d: number[] = Array(n + 1).fill(0);
  let s = 0;
  for (let i = 0; i < n; ++i) {
    s += d[i];
    nums[i] += s;
    if (nums[i] === 0) {
      continue;
    }
    if (nums[i] < 0 || i + k > n) {
      return false;
    }
    s -= nums[i];
    d[i + k] += nums[i];
  }
  return true;
}