クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「3-2 4 Values whose Sum is 0(POJ No.2785)」(p.147)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「3-2 4 Values whose Sum is 0(POJ No.2785)」(p.147)
のコードをPythonで書き直したものとなっています。

入力

n\\A_1\,…\,A_n\\B_1\,…\,B_n\\C_1\,…\,C_n\\D_1\,…\,D_n

入力例



6
-45 -41 -36 -36 26 -32
22 -27 53 30 -38 -54
42 56 -37 -75 -10 -6
-16 30 77 -46 62 45


解答

from bisect import bisect_left, bisect_right

# 入力
n = int(input())
A = list(map(int,input().split()))
B = list(map(int,input().split()))
C = list(map(int,input().split()))
D = list(map(int,input().split()))

CD = [0] * (n * n)    # CとDの組からの取り出し方

# CとDからの取り出し方の組を全列挙
for i in range(n):
    for j in range(n):
        CD[i * n + j] = C[i] + D[j]

CD.sort()

res = 0
for i in range(n):
    for j in range(n):
        cd = -(A[i] + B[j])
        # CとDから和がcdとなるように取り出す
        res += bisect_right(CD, cd) - bisect_left(CD, cd)

print(res)