LeetCode 974. 和可被 K 整除的子数组

发表于 更新于 分类于 Changyan： 本文字数： 427 阅读时长 ≈ 2 分钟

题目描述

和可被 K 整除的子数组

给定一个整数数组 A，返回其中元素之和可被 K 整除的（连续、非空）子数组的数目。

示例 1 :

输入：A = [4,5,0,-2,-3,1], K = 5
输出：7
解释：
有 7 个子数组满足其元素之和可被 K = 5 整除：
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

解法一（哈希表）

看着题目与 LeetCode 560. 和为K的子数组 很像，就依葫芦画瓢地用了hash表的方法，时间复杂度为O(n2)，想法很好，不幸的是超时了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
public:
  int subarraysDivByK(vector<int>& A, int K) {
    unordered_map<int, int> table;
    table.insert({0, 1});
    int sum = 0, ans = 0;
    for (int i = 0; i < A.size(); i++) {
      sum += A[i];
      auto iter = table.begin();
      while (iter != table.end()) {
        if ((sum - iter->first) % K == 0)
          ans += iter->second;
        iter++;
      }
      if (table.find(sum) != table.end())
        table.find(sum)->second++;
      else
        table.insert({sum, 1});
    }
    return ans;
  }
};

解法二（取余）

其实两道题确实有关联，都是前缀和的题目，不过用hash表来记录前缀和的方法不适合这道题。

考虑到如果两个前缀和对 K 进行取余的结果相等，则这两个前缀和之间的元素和可以被 K 整除，由于我们可以用大小为K的数组去记录所有前缀和的余数分布，然后累加起来得到最终结果。

取余公式：(x % K + K) % K

还要记得额外处理前缀和对 K 取余为0的情况。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution {
 public:
  int subarraysDivByK(vector<int>& A, int K) {
    vector<int> count(K);
    int sum = 0;
    for (int i = 0; i < A.size(); i++) {
      sum += A[i];
      count[(sum % K + K) % K]++;
    }
    int ans = count[0];
    for (int i = 0; i < K; i++) {
      ans += count[i] * (count[i] - 1) / 2;
    }
    return ans;
  }
};