AtCoder ABC 344 C - A+B+C (5Q, 灰色, 250 点)

とても教育的な問題！　ちゃんと解くには、計算量の理解が必要となる問題である。

• 問題へのリンク

問題概要

3 つの数列

• 
• 
• 

が与えられる。これらの数列に対して、次の 個のクエリに答えよ。

【クエリ】
整数 が与えられるので、 からそれぞれ 1 個ずつ選んで和をとると に一致させられるかどうかを判定せよ。

制約

• 
• 
• 

考えたこと

単発のクエリならば、次のように全探索する解法も考えられる。

bool res = false;
for (int i = 0; i < N; i++) {
    for (int j = 0; j < M; j++) {
        for (int k = 0; k < L; k++) {
            if (A[i] + B[j] + C[k] == X) res = true;
        }
    }
}

しかし、この解法を 回のクエリすべてに実行すると、計算量は となる。これは TLE してしまう。

ここで、多くの言語で用意されているset 型は、通常、

• 要素を挿入する操作
• 要素を検索する操作

を高速に実行できることを思い出そう。C++ の set 型の場合、挿入・検索にはそれぞれ の計算量で実行できる（要素数を としている）。

今回の問題においても、次のようにすることが考えられる。

---

• あらかじめ、数列 から 1 個ずつ選んで和をとってできる数値をすべて集合 に挿入しておく
• 各クエリについて、値 が集合 に含まれるかどうかを判定する

---

C++ の set 型を用いた場合の計算量を評価しよう。前者は、集合 に挿入する操作を 回実行するので、計算量は となる。

次に、 回のクエリについて、値 が集合 に含まれるかどうかを判定する処理の計算量は となる。

合わせると、計算量は となる。これならば AC をとれる。

コード

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

int main() {
    int N, M, L, Q;
    cin >> N;
    vector<int> A(N);
    for (int i = 0; i < N; i++) cin >> A[i];
    cin >> M;
    vector<int> B(M);
    for (int i = 0; i < M; i++) cin >> B[i];
    cin >> L;
    vector<int> C(L);
    for (int i = 0; i < L; i++) cin >> C[i];

    set<int> S;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < M; j++) {
            for (int k = 0; k < L; k++) {
                S.insert(A[i] + B[j] + C[k]);
            }
        }
    }

    cin >> Q;
    for (int _ = 0; _ < Q; _++) {
        int X;
        cin >> X;
        if (S.count(X)) cout << "Yes" << endl;
        else cout << "No" << endl;
    }
}