Codeforces Round #707 Div1A. Going Home 畳み込みを使った作れる和の数

畳み込みを使って解きます。

畳み込みで、和を求めるという演算をどうすればいいかわかりませんでしたが、考えてみるとシンプルでした。ただし、自分自身を畳み込みんでいることや、重複を排除して演算するときの配慮が必要です。

題意

• n個の数字が与えられる。 ($n \leq 2 \cdot 10^{5}$)
• 要素は$1 \leq a_i \leq 2500000 $である
• すべて異なるインデックス$i, j, k, l$に対して、$a_i + a_j = a_k + a_l$が存在するか？存在するなら1つ例を挙げよ

例

• パターン1: [1,2,2,4,5,7]のとき、例えば $1+5 = 2+4$を作れます(すべて別の数字)ほかに$2+7 = 4+5$でもよいです。
• パターン2: [1,2,2,3]のとき、$1+3 = 2+2$です。これはパターン2に似ていますが少し異なります。
• パターン3: [1,1,2,2,10,100]のとき、$1+2 = 1+2$を作れます(ペアが2つ)
• パターン4: [1,1,1,1,10,100]のとき、$1+1 = 1+1$を作れます(ある数字が4つ)
• NG: [1,3,10,100]や[2,2,2,3]は満たすものを作れません。NOです。

こう考えた

まず、各数を操作し、それぞれの数($1から2,500,000まで$)の出現回数を数え上げます。数$x$に対する出現回数を$cnt(x)$とします。この操作回数のテーブルを線形走査します。ある$x$について

• $cnt(x) \leq 4$のとき、xを4つ返せばよいです
• $cnt(x) \leq 2$となるxが2つ存在(例えば$x_1, x_2$とする)するとき、$x_1 + x_2 = x_1 + x_2$ですからそれを返せばよいです。

これによりパターン3,4が存在するかを判定できます。それでは、パターン1,2が存在するかを求めます。

cntを使って、別の数列bを作ります。$b_i = max(1, cnt(i))$とします。つまり、$b_i$は$i$が数列$a$に１度でも出現すれば1, 出現しなければ0のテーブルです。これを自分自身で畳み込みます。結果を数列cとして、$c = convolution(b, b)$を計算します。数列$b$の特性を考えると、数列cは長さ$5,000,001$であり、$c_k$とは、$b_i = b_k = 1$である数字$i,j$を選んで何個作れるかを示します。この時、畳み込みの結果には、$i=j$、$ij$のときが含まれていることに注意します。

• $c_k$ = 0の時、$k$は作ることができません。
• $c_k$ = 1のとき、$k$は$i=j$と選んで、1つだけ作ることができます。
• $c_k = 2$の時、$k$は$i \neq j$な数字i,jのペアを使って1つだけ作ることができます$c_k=2$になるのは、$i < j$と$i > j$の時の2回が含まれたからです。
• $c_k=3$の時、は考察が必要です。例えば、[1,2,3]というとき、$c_4 = 3 $になります。$1+3, 2+2, 3+1$が含まれるからです。ただし、数列bは$a_i$の重複を考慮していない(複数でも1にしている)ため、元の数列$a$が[1,2,3]なのか[1,2,2,3]なのかは区別できません。前者は条件を満たせませんが、後者は例に挙げたパターン2のとおり満たします。これを確認するため、$k$が偶数の場合、$cnt(k) >= 2$であれば、条件を満たします。$k$が奇数の場合はこれを満たす条件はありません。
• $c_k \geq 4$の時、その数字はすべて異なる適当な最低4つの数字で作ることができます。

こうすると、$c_k=4$あるいの時か、$c_k=3$であり条件を満たすときは$k$を作れることがわかりました。
$i=1$から順に$cnt(i) \leq 0$かつ$cnt(k-i) \leq 0$な数を探し、それを出力すればよいです。

補足: 考えたことの図示

• 左の図を見ます。まず、入力[1,2,2,4,5]を出現回数ことにマッピングします。1 indexとするなら、[1,2,0,1,1]です。これの1度でも出現している数字は1とする数列$b$を作ります。[1,1,0,1,1]となります。この数列b同時を畳み込みます。[0,1,2,1,2,4,2,1,2,1]となりました。$c_6=4$ですが、これは下のように$1+5,2+4,4+2,5+1$の4つがカウントされたためです。これを出力します。

• 右の図を考えます。入力[1,2,2,3]の場合、2は2回出現したことは覚えておき、数列bを作ります。結果、[0,1,2,3,1,1]となります。$c_4=3$ですが、これは$1+3, 2+2, 3+1$がカウントされた結果です。$2+2$が2回出現できないかを確認するため、$cnt(2)$を調べると2です。このため、$c_k$は作れます。

補足: うまくいかない例

今回、わざわざ、数列bを作り畳み込みました。果たして、出現回数を1に制限する必要があるのでしょうか？回数をそのまま畳み込んでみます。

- 左の図はうまくいくように見えます。[0,1,4,4,2,6,4,1,2,1]となるため、なんとなく、最大の6を値使うとうまくいきそうです。ところが、$c_k \leq 4$であるほかの$k$はどうでしょうか？$k=3,4,7%などは条件を満たせるようには選べません。
- 右の図も同様です。[2,2,2,3]のとき、$c$は[0,0,0,9,6,1]になります。が、同様に条件を満たせる数字は選べません。
畳み込みを行うと、$a_i$と$b_j$からの情報を使った$c_k$には($k=i+j)一切$i,j$の譲歩が残りません。このことに注意する必要があります。

実装

AtCoderのACLを使った例を示します。問題では、indexを答える必要があったため、最初にindexを保存しておき出力しています。(1-indexに注意)

#define FOR(i, begin, end) for(int i=(begin),i##_end_=(end);i<i##_end_;i++)
#define REP(i, n) FOR(i,0,n)
using namespace std;
using namespace atcoder;
#define SZ 2500100
#define ll long long

int main(int argc, char *argv[]) {
    int n; cin >> n;
    ll x;
    int prev2 = -1;
    vector<long long> a(n);
    vector<long long> freq(SZ+10);
    vector<long long> b(SZ+10);
    vector<vector<ll>> index(SZ+10);
    REP(i, n) {
        cin >> x;
        index.at(x).emplace_back(i + 1);
        freq.at(x)++;
    }
    /* 畳み込みを使わないで十分なパターン */
    REP(i, SZ+1){
        /* ある数字が4つ以上あるならx+x+x+x */
        if(freq.at(i) >= 4) {
            cout << "YES" << "\n" << index.at(i)[0]<< " " << index.at(i)[1]<< " " << index.at(i)[2]<<" " <<  index.at(i)[3]<<"\n";
            return 0;
        }
        /* ある数字が2つ以上、が2組あるならx+y = x+y */
        if(freq.at(i) >= 2) {
            if(prev2 == -1){ prev2 = i; continue; }
            cout << "YES" << "\n"<< index.at(i)[0]<< " " << index.at(prev2)[0]<< " " << index.at(i)[1]<<" " <<  index.at(prev2)[1]<<"\n";
            return 0;
        }
    }
    /* 畳み込みをする際、いったん、各数字は1度しか出現しないことにする */
    REP(i, SZ+1) {if(freq.at(i) > 0) b.at(i) = 1; }
    vector<long long> c = convolution(b, b);
    long long target = -1;
    REP(i, 2*SZ+1) {
        /* a + a = x + y のパターン
         * この数を2aとすると、
         * aが2回以上出現しても1回にしてしまっている。このため
         * この結果が2になるのは、x+yが存在するとき。
         * ならば、aが2つ以上あるなら、a + a = x + yは成立する。*/
        if(c.at(i) >= 2 && ((i&1) == 0) && freq.at(i/2) >= 2) {target = i; break; }
        /* 畳み込んで、4以上ならば、それを満たす数x,y の組は2つ以上絶対にある */
        if(c.at(i) >= 4){target = i; break;}
    }
    /* 数がない場合はNO */
    if(target == -1) {cout << "NO" << "\n"; return 0;}
    int foundCnt = 0;
    vector<ll> res;
    /* targetを作れる組を2つ見つける */
    FOR(i, 1, target){
        if(freq[i] > 0 && freq[target - i] > 0){
            res.emplace_back(index.at(i).front());
            res.emplace_back(index.at(target - i).back());
            foundCnt++;
        }
        if(foundCnt == 2){
            cout << "YES\n" << res[0] << " " << res[1] << " " << res[2] <<  " " << res[3] << "\n";
            return 0;
        }
    }
    /* NO REACH */
    return 0;
}