[python] 백준 2143 : 두 배열의 합

문제 링크
https://www.acmicpc.net/problem/2143

접근

구간합과 딕셔너리를 적절히 사용하면 되겠다 생각했다.

Asum = [0 for _ in range(n + 1)]
Bsum = [0 for _ in range(m + 1)]

for i in range(n):
    Asum[i] = Asum[i - 1] + A[i]

for i in range(m):
    Bsum[i] = Bsum[i - 1] + B[i]

배열 A와 B 의 0부터 i번 까지의 합을 미리 구해놓는다. 편의를 위해 n+1, m+1개를 할당해놓고, Asum[-1] 또는 Bsum[-1]을 0으로 한다.

Ad, Bd = {}, {}
for i in range(n):
    for j in range(i, n):
        x = Asum[j] - Asum[i - 1]
        Ad.setdefault(x, 0)
        Ad[x] += 1

for i in range(m):
    for j in range(i, m):
        x = Bsum[j] - Bsum[i - 1]
        Bd.setdefault(x, 0)
        Bd[x] += 1

그리고 이제 딕셔너리를 이용하여 구간합의 값이 얼마나 많이 나왔는지 저장한다.
예를 들어, A[1] 부터 A[5] 까지의 합이 10이 나왔으면 Ad[10] 을 1 증가시킨다.
즉 Ad[x] = 구간합이 x가 나온 것의 개수이다. (Bd도 마찬가지)

ans = 0
for x, cnt in Ad.items():
    if T - x in Bd:
        ans += cnt * Bd[T - x]

마지막으로 딕셔너리 Ad의 key 값 x에 대하여 Bd[T-x] 가 존재하면, Ad[x] * Bd[T-x] 를 답에 더해주는 식으로 최종 답안을 구할 수 있다(여기서는 Ad[x] 가 cnt이다).

정답 코드

import sys

sys.setrecursionlimit(10 ** 8)
input = lambda: sys.stdin.readline().rstrip()

T = int(input())

n = int(input())
A = list(map(int, input().split()))

m = int(input())
B = list(map(int, input().split()))

Asum = [0 for _ in range(n + 1)]
Bsum = [0 for _ in range(m + 1)]

for i in range(n):
    Asum[i] = Asum[i - 1] + A[i]

for i in range(m):
    Bsum[i] = Bsum[i - 1] + B[i]
    
Ad, Bd = {}, {}
for i in range(n):
    for j in range(i, n):
        x = Asum[j] - Asum[i - 1]
        Ad.setdefault(x, 0)
        Ad[x] += 1

for i in range(m):
    for j in range(i, m):
        x = Bsum[j] - Bsum[i - 1]
        Bd.setdefault(x, 0)
        Bd[x] += 1

ans = 0
for x, cnt in Ad.items():
    if T - x in Bd:
        ans += cnt * Bd[T - x]
print(ans)