E. Special Elements (Codeforces Round #640 Div. 4)

制約について考えさせられる問題でした。

問題

https://codeforces.com/contest/1352/problem/E

問題概要

長さnの数列が与えられる。各項は1以上n以下の整数である。

各項の値について、その数列の連続する2つ以上の項の和で表すことができる場合、その項は「特別」であるとする。

与えられた数列に「特別」な項がいくつあるかを答えよ。

考察

TLが1秒でMLが64MBと、制約がきついです。

私はきっちりTLEもMLEも取りました。とほほ。

時間計算量

まず、累積和を取って一つ一つしゃくとり法で確かめていくという実装をしたところ、TLEでした。（しゃくとり法はO(n)で全体でO(n²)だから間に合うと思ったのですが・・・。最悪ケースで左右のポインタが一つずつ交互に動くケースが思ったより遅いということのようです。落ちたテストケースを参照してください。）

https://codeforces.com/contest/1352/submission/79522413

空間計算量

和の探索に時間がかかるなら前もってすべて列挙してしまえばいいのでは？と思って実行に移してはいけません。この問題はMLが64MBです。

https://codeforces.com/contest/1352/submission/79529811

はい、見事にMLEで落ちます。

最終的に

最初の方針でしゃくとり法を二分探索に変えると、割とギリギリですがなんとか通ります。

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

int main() {
    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        vector<int> a(n);
        vector<int> sum(n + 1);
        for (int i = 0; i < n; i++) {
            cin >> a[i];
            sum[i + 1] = sum[i] + a[i];
        }
        int ans = 0;
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n - 1; j++) {
                auto itr = lower_bound(sum.begin() + j + 2, sum.end(), a[i] + sum[j]);
                if (itr != sum.end() and *itr == a[i] + sum[j]) {
                    ans++;
                    break;
                }
            }
        }
        cout << ans << endl;
    }
}

writerの想定解

Editorialによれば

• まず値の「出現数」をカウントする
• 2重ループを使って連続する項の和をすべて試す
  • 和の値がまだ出現していなければ、最初に集計した出現数分答えに足す
    • 二度以上足さないように注意
  • 条件より和がnを超えた場合はそれ以上計算する必要がないので、そこで打ち切り
  • 逐次計算するので、メモリに和を溜め込む必要がない

ということだそうです。

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

int main() {
    int t;
    cin >> t;
    while (t--) {
        int n;
        cin >> n;
        vector<int> a(n);
        vector<int> cnt(n + 1);
        for (int i = 0; i < n; i++) {
            cin >> a[i];
            cnt[a[i]]++;
        }

        int ans = 0;
        for (int i = 0; i < n - 1; i++) {
            int sum = 0;
            for (int j = i; j < n; j++) {
                sum += a[j];
                // i == j のときはまだ一つしか足していないので続行
                if (i == j) continue;
                if (sum <= n) {
                    ans += cnt[sum];
                    // 同じ値を二度以上足さないように
                    cnt[sum] = 0;
                }
            }
        }
        cout << ans << endl;
    }
}

私が通した回答より10倍以上速いです。（白目