AGC023: A - Zero-Sum Ranges

問題
思考過程
解法
知見
感想
参考

これ 200 点なの？という問題だった．難しい．

問題

AGC023: A - Zero-Sum Ranges

f:id:pytwbf201830:20190226233159p:plain

思考過程

制約から，

O(N²) は無理．O(N) or O(NlogN) 以下で回す必要がある．
全探索は厳しそう．
累積和を使った区間の和を利用，計算するのでは？
ただ，全部の場合を数えなくて良いと考えられる（全探索は厳しいため）．他の方法で求められる．

など基本的なことを考えていた．「累積和を使った区間の和」のイメージは，以下のような感じ．

部分列を何個足すか：部分列の選び方
N   個の和  ： N -  (N)  + 1 = 1 通り
N-1 個の和  ： N - (N-1) + 1 = 2 通り
N-2 個の和  ： N - (N-2) + 1 = 3 通り
...
3 個の和    ： N - (3) + 1   = N-2 通り
2 個の和    ： N - (2) + 1   = N-1 通り
1 個の和    ： N - (1) + 1   = N   通り

上記のように区間の選び方の全探索を考えた．しかしこれだと， $N = 2*10^{5}$ のときの最悪時間計算量が $10^{10}$ となってしまい TLE（ $10^{5}$ を探索するとき区間の取り方が $10^{5}$ 通りあるため）．結局この全探索の解法しか書けなかった．

念のため，以下に TLE したコードを載せておく（アルメリアさんの記事で，解法が分からなくても「愚直解，全探索を書けることはとても大事だ」と言う記述を見たので）．

▶︎ソースコード展開

#include <iostream>
#include <string>
#include <vector>
#include <iterator>
#include <algorithm>
#include <functional>  // std::greater<T>()
#include <utility>  // std::swap()
#include <numeric>  // accumulate(ALL(vec), 0)  0 は初期値
using namespace std;
using ll = long long;

#define REP(i, n) for (ll i=0; i < n; i++)  // 0 ~ n-1

int main() {
    // input
    ll N;
    cin >> N;
    vector<ll> vec(N);
    REP(i, N) cin >> vec[i];
    
    // calculation
    if (N == 1) {
        if (vec[0] == 0) cout << 1 << "\n";
        else cout << 0 << "\n";
    } else if (N == 2){
        ll res = 0;
        if (vec[0] == 0) res++;
        if (vec[1] == 0) res++;
        if (vec[0] + vec[1] == 0) res++;
        cout << res << "\n";
    } else {
        // 累積和
        vector<ll> partial = {0}; // 累積和を取るための 0
        copy(vec.begin(), vec.end(), back_inserter(partial));
        partial_sum(partial.begin(), partial.end(), partial.begin());
        
        // 単独で 0 の場合
        ll res = 0;
        for (ll i = 0; i < N; ++i) {
            if (vec[i] == 0) res++; 
        }
        
        // 2 個以上で 0
        ll left, right, interval_sum;
        for (ll i = N; i >= 2; --i) {
            for (ll j = 1; j <= N-i+1; ++j) {
                left = j-1, right = j+i-1;
                interval_sum = partial[right] - partial[left];
                if (interval_sum == 0) res++;
            }
        }
        
        cout << res << "\n";
    }
    
    return 0;
}

解法

解説が短かったのでそのまま載せます．

S を長さ N + 1 の数列とし、S0 = 0, Si = Si−1 + Ai とします。ある連続する部分列の総和は、S の 2 要素の差です。よって、和が 0 になる連続する部分列の数は、S の中で値の等しい 2 つの要素を選ぶ方法の数に等しいです。これは、S の要素を sort するなどして、同じ値の個数を数えることで計算出来ます。sort に O(NlogN) かかるため、この問題は O(NlogN) で解けました。

整理すると，

先頭に 0 を付け加えた累積和の数列 S を用意．
i ~ j (i < j) までの部分列の総和 $S_{i~j}$ は，S[j] - S[i-1] で求められる（累積和のテクニック）．
連続する部分列の和が 0 になるとき，S[j] == S[i-1] となるため，数列 S で同じ数のペアを探せば良い．

となる．あとは上記を実装するのみ．計算量は O(NlogN)．

#include <iostream>
#include <vector>
#include <map>
#include <algorithm>
using namespace std;
using ll = long long;

int main() {
    // input
    ll N;
    cin >> N;
    vector<ll> A(N), S(N+1, 0);
    
    for (ll i = 0; i < N; ++i) {
        cin >> A[i];
        S[i+1] = S[i] + A[i];  // 累積和
    }
    
    // calculation
    ll res = 0;
    map<ll, ll> mp;
    for (ll i = 0; i < N+1; ++i) {
        res += mp[S[i]];  // key がなければ 0 を返す
        mp[S[i]]++;
    }
    
    cout << res << "\n";
    
    return 0;
}

実装のポイントは二つ挙げられる，

i より右で S[i] と等しいものの数を数える
std::map で現れた頻度を持ちながら組み合わせの数を数えていく

二つ目が少し分かりにくいので説明を加える．

i 番目までにその数が k 個現れていれば，S[i]（i 番目の数）と，それより前に存在する k 個の数との組み合わせが k 個作れる．上記のコードでは，それを std::map で保持しながら組み合わせを数えている（ちなみに，std::map は key が存在しない場合はエラーではなく 0 を返す．それも込みで実装されているので注意）．

以上の仕組みが分かっていれば，この問題は，配列を使わずとも簡潔に実装することができる．見通しの良いコードになっている．

#include <iostream>
#include <map>
using namespace std;
using ll = long long;

int main() {
    // input
    ll N, A_i, sum = 0, res = 0;
    cin >> N;
    map<ll, ll> mp;
    mp[0] = 1;
    
    for (ll i = 0; i < N; ++i) {
        cin >> A_i;
        sum += A_i;  // 累積和の計算
        res += mp[sum];  // key がなければ 0 を返す
        mp[sum]++;
    }
    
    cout << res << "\n";
    
    return 0;
}

知見

累積和を取るときは，配列 A のサイズを N+1 で取り A[0] = 0 とした上で累積和を取ると区間の和を O(1) で算出できる．以下を参考にさせていただいた．

qiita.com

感想

解法が頭良すぎてすごい．累積和を使って区間の和を求めるという方針はそこまで悪くなかったのかなとも思うが（本番で解説の解法を思いつけと言われたら無理），恐らく区間の和を考えるまではほとんどの人が辿り着いている気がする．そこから解けるかどうかが分かれ目だったと考えられる．引き続き精進頑張ります．

参考

ferin-tech.hatenablog.com

pyてよn日記

一寸先は闇が人生

AGC023: A - Zero-Sum Ranges

問題

思考過程

解法

知見

感想

参考