776-C. Molly's Chemicals

2021-03-27 20:18:58 +0800 CST Tags: prefix-sum, sequence

出處：https://codeforces.com/contest/776/problem/C
提交：https://codeforces.com/contest/776/submission/111320970

問題：給定 $n\in[1,10^5]$ 個介於 $[-10^9,10^9]$ 的整數，以及常數整數 k 其絕對值介於 $[1,10]$ 之間，請問總共有幾個連續子區間的和是為 k 的非負整數次方？

解法：我們很容易知道這題的次方數其實可以一個一個依序嘗試，只要目標和超過全體最大可能和或者小於全體最小可能和就可以停下來然後注意一下 $k=\pm1$ 的特殊情況就好。這題最困難的點應該在於怎麼把一個陣列中所有總和為指定數字的連續子區間數量求出來，而這道子題其實正好就是 LeetCode 560. Subarray Sum Equals K [1]，它是有已知最佳解的，官方解法很巧妙地利用了 map [2] 的特性，讓我們在不確定 key 的情況之下還能每次都在 $\mathcal O(\log_2{n})$ 的時間下求值。精神如下，所有的子區間都可以被歸為其中一個以某特定元素為結尾的類別，舉例來說，長度為 3 的陣列 arr 的所有子區間就可以分類為 C$_1$ (以 arr[1] 結尾)、C$_2$ (以 arr[2] 結尾)，以及 C$_3$ (以 arr[3] 結尾)，那對於每個分類 C$_i$，假設目標後綴和是為 k，但從頭到現在的前綴和 arr[1] + … + arr[i] 是為 sum，那麼 C$_i$ 的種類數其實就是到目前為止的前綴和為 sum-k 的種類數，因為前綴和與後綴和會互補，所以我只要另外維護 prefix sum 跟從頭到現在的前綴和為 k' 的種類數 map[k'] 即可，然後記得一開始要 map[0] = 1。

附註：
[1] https://leetcode.com/problems/subarray-sum-equals-k/
[2] 這邊如果不小心選成 unordered_map 的話會造成 TLE，因為兩者的內部實作結構不同，map 是平衡樹而 unordered_map 是雜湊表，後者最差會到 $\mathcal O(n)$ 而前者則是 $\mathcal O(\log_2{n})$，故一般競程比賽下通常會選擇前者。

實作：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

#include <bits/stdc++.h>
using namespace std;

int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); // IO 優化
    long long target=1, ans=0, sum;
    map<long long, long long> count;
    int n, k, arr[100000]; cin >> n >> k;
    for (int i=0; i<n; i++) cin >> arr[i];
    do {
        sum = 0; count.clear(); count[0] = 1;
        for (int i=0; i<n; i++) {
            ans += count[sum - (target-arr[i])];
            sum += arr[i];
            count[sum]++;
        }
        target *= k;
    } while (k!=1 && !(k==-1 && target==1) && abs(target)<=1e14);
    cout << ans;
    return 0;
}

結果：