ABC344 C問題の解説(Python)

はじめに

AtCoder Beginner Contest(略称:ABC)とは、AtCoderが開催している、競技プログラミングコンテストのことです。

今回はABC344のC問題を解説していきます。

ABC344 C – A+B+C

問題

問題

3個の整数列 \(A=(A_1, …,A_N ), B=(B_1, …,B_M ), C=(C_1, …,C_L )\) と整数列 \(X=(X_1, …,X_Q )\) が与えられます。Xの各整数について、A, B, C から 1 つずつ整数を選んで和をとることで、同じ値にできるかを調べる問題です。

制約

  • \(1 \leqq N,M,L \leqq 100\)
  • \(0 \leqq A_i,B_i,C_i \leqq 10^8\)
  • \(1 \leqq Q \leqq 2×10^5\)
  • \(0 \leqq X_i \leqq 3×10^8\)

思考の筋道

A, B, Cの要素数は最大100で、Xの要素数は最大\(2×10^5\)です。そのため、A, B, C, X に関する4重の全探索を行うとTLEです。

まずはA, B, Cだけを見るようにしましょう。事前に3つの整数の和をすべて求め、集合に保管することを考えます。A, B, Cの3重の全探索であれば、計算量は10^6なので間に合います。

その後に、Xの各整数が集合に属しているかどうかを調べましょう。集合に要素が属するかどうかは、平均計算量 \(O(1)\) で調べられるので、Xの各整数を調べても十分間に合います。

コード

N = int(input())
A = list(map(int, input().split()))
M = int(input())
B = list(map(int, input().split()))
L = int(input())
C = list(map(int, input().split()))
Q = int(input())
X = list(map(int, input().split()))

S = set()

for num_A in A:
    for num_B in B:
        for num_C in C:
            S.add(num_A + num_B + num_C)

for num_X in X:
    if num_X in S:
        print("Yes")
    else:
        print("No")

リンク