🔗 🟡 974. Subarray Sums Divisible by K 1676

tags: `Weekly Contest 119` `前綴和(Prefix Sum)` `雜湊表(Hash Table)` `同餘定理`

題意

給定一個整數陣列 $nums$ 和一個整數 $k$ ，返回具有可以被 $k$ 整除的和的 非空子陣列 的數量。

子陣列 是陣列的連續部分。

約束條件

$1 \leq \text{nums.length} \leq 3 \cdot 10^4$
$-10^4 \leq \text{nums}[i] \leq 10^4$
$2 \leq k \leq 10^4$

思路：前綴和(Prefix Sum) + 雜湊表(Hash Table) + 同餘定理

我們可以使用前綴和(Prefix Sum) 來計算陣列中每個位置的前綴和，對於子陣列 $[l, r]$ 的和可以表示為 $s[r] - s[l-1]$ ，若其為 $k$ 的倍數，則 $s[r] \equiv s[l-1] \pmod{k}$ 。

故可以記錄每個前綴和對 $k$ 取模後的結果，並使用一個雜湊表或長度為 $k$ 的陣列 $cnt$ 來記錄取模後的值出現的次數。當枚舉到第 $i$ 個位置時，我們可以計算前 $i$ 個位置的前綴和 $s$ ，並對 $k$ 取模，則 $cnt[s \bmod k]$ 代表前面有多少個位置與當前位置的前綴和對 $k$ 取模後的結果相同，故能組成需要的子陣列，可以累加其出現次數至答案中。

具體步驟如下：

使用一個雜湊表或長度為 $k$ 的前綴和陣列 $cnt$ 來記錄每個前綴和餘數出現的次數，初始化時設置 $cnt[0] = 1$ ，代表空陣列。
遍歷陣列 $nums$ ，對每一個元素累加計算前綴和 $s$ 。
計算當前前綴和 $s$ $s$ 對 $k$ $k$ 的餘數 $s \bmod k$ $s mod k$ 。
- 由於我們只在意餘數，因此可以每次都將 $s$ 對 $k$ 取模。
- 此外，由於約束條件中 $nums[i]$ 可以為負數，而在 C++ 中，對負數取模的結果可能為負數，因此我們需要對取模後的結果加上 $k$ 並再次對 $k$ 取模，以確保結果為正數。
將當前餘數出現的次數累加到答案中，這表示以當前元素為結尾的子陣列中，有多少個子陣列是可以被 $k$ 整除的。
更新雜湊表 $cnt$ 中當前餘數出現的次數。

複雜度分析

時間複雜度： $\mathcal{O}(n)$ ，其中 $n$ 為陣列的長度。
空間複雜度： $\mathcal{O}(min(n, k))$ $O (min (n, k))$ ，雜湊表的空間複雜度。
- 若使用長度為 $k$ 的陣列 $cnt$ 來記錄，則空間複雜度為 $\mathcal{O}(k)$ 。

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        ans = 0
        s = 0 # prefix sum
        cnt = defaultdict(int)
        cnt[0] = 1
        for x in nums:
            s += x
            ans += cnt[s % k]
            cnt[s % k] += 1
        return ans

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        int ans = 0, s = 0;
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        for (int x : nums) {
            s = (s + (x % k) + k) % k; // 避免出現負數
            ans += cnt[s];
            cnt[s]++;
        }
        return ans;
    }
};