Codeforces Round #632 (Div. 2) - C

問題へのリンク

C - Eugene and an array
codeforces.com

問題概要

長さ の数列 が与えられる。区間の和が となる連続部分列を含まない連続部分列の個数を求めよ。

制約

考察

ある条件を満たす数え上げの問題である。
本番では、さっと順序立てて書けずに、こんがらがって時間を溶かしまくったので、反省として記事を書く。

まず累積和を取る。
例のごとく、
左から 個目まで取った時の累積和
としておく。
このとき、区間 の和は、 で求められる。
すると、ある区間 の和が になるということは、 を用いて以下のように表現できる。

つまり

(この段落の説明は、後述で図でも示している。)
出力する答えを 、今見ている連続部分列を としておく。
基本的な方針として、 を伸ばしていくのだが、新たに追加する要素 までの累積和 の値が、任意の の要素までの累積和と異なるならば、そのまま伸ばすことができる。
問題は がある場合である。
このときはまず、 の中で取れる連続部分列の数を に加算しておく。
このときに加算される値は、 の長さを とすると、 と計算できる。
の左端を としておき、 がなくなるまで、 を右に動かしていく。
これは、 に含まれる要素までの累積和を set などで持っておくと、尺取り法の要領でできる。
ここまでの処理をしてから、連続部分列 を伸ばす。
右端が までいったときにも、最後に を に加算する。

ただ、これだけでは不十分である。
というのも、重複が起こり得るからで、その分を取り除いておく必要がある。
例えば、 であるような場合を考えてみる。
上述の手順で実験する。



ここまでは難なく進むが、問題は次の場合 () である。

連続部分列に を入れようとしているが、 と等しい累積和が含まれているので、まだ含めることができない。

まず、上図の矢印で囲まれた で取れる連続部分列の個数を に加算しておく。

次に、上図のように累積和の等しい値がなくなるまで左端を動かす。

ここまできてようやく連続部分列に を含めることができる。

ただ、注意深く見ると分かるが、上図の矢印部分が重複して数え上げられる。
この部分は、後の処理で再び加算されるので、 を含める前に、この連続部分列で取れる連続部分列の個数を から差し引いておく。

これで問題が解決できる。

本番では、どのデータ構造が適切かと考えているときに、重複をどう処理すればいいのかという考えに至り、ものすごくこんがらがっていた。
考えを纏めて素早くコードを書けるようになりたい。

コード

本番では尺取り部分を while 文で書いていたが、for 文で書いておく。

#include <iostream>
#include <vector>
#include <set>
using namespace std;

int main(){
    int n;
    cin >> n;
    vector<int> a(n);
    vector<int64_t> S(n + 1);
    S[0] = 0;
    for(int i = 0; i < n; ++i){
        cin >> a[i];
        S[i + 1] = S[i] + a[i];
    }
    
    int64_t ans = 0;
    int left = 0;       // 今見ている連続部分列の左端
    int len = 0;        // 今見ている連続部分列の長さ
    set<int64_t> st;    // 累積和に関する集合
    st.insert(0);
    for(int right = 1; right <= n; ++right){
        // S[r] が st に含まれていなければ、そのまま連続部分列を伸ばす
        if(st.find(S[right]) == st.end()){
            st.insert(S[right]);
            ++len;
        }
        // S[r] が st に含まれているとき
        else{
            // 直前までの連続部分列で取れる連続部分列の個数を加算
            ans += (int64_t)len * (len + 1) / 2;
            // S[r] がなくなるまで左端を動かす
            while(st.find(S[right]) != st.end()){
                st.erase(S[left++]);
                --len;
            }
            // 後にも数え上げられる重複分を取り除く
            ans -= (int64_t)len * (len + 1) / 2;
            // 右端を伸ばす
            st.insert(S[right]);
            ++len;
        }
    }
    // 残る連続部分列の分を加算
    ans += (int64_t)len * (len + 1) / 2;
    cout << ans << endl;
}

雑記

記事を書くのも久々になってしまった。
相変わらず冗長すぎる。

KMP なり学んだアルゴリズムを書こうと思っていたが、めんどくさくなって下書きに溜まっていく一方である。
ブログを備忘録として用いるために、ちゃんと書こうね…

Atcoder では緑から水 (現在 : 緑)、Codeforces では緑から青 (現在 : 青) をさまよっている。
今回で再び青になった。
全く安定しなくて笑うほかない。

追記

この操作を使える問題には、例えばこれがある。
D - Xor Sum 2