bj2143 두 배열의 합

문제

• 백준 URL: 두 배열의 합
• 난이도 골드 3

풀이

import sys
from collections import defaultdict

input = sys.stdin.readline
T = int(input())
n = int(input())
A = list(map(int,input().split()))
m = int(input())
B = list(map(int,input().split()))

sum = defaultdict(int)
for i in range(n):
    s_a = 0
    for j in range(i,n):
        s_a += A[j]
        sum[s_a] += 1

ans = 0
for i in range(m):
    s_b = 0
    for j in range(i,m):
        s_b += B[j]
        r = T - s_b
        if r in sum:
            ans += sum[r]

print(ans)

이 문제는 두 배열 A와 B의 모든 부 배열(연속된 부분 배열)의 합을 각각 구한 뒤, A의 부 배열 합과 B의 부 배열 합을 더해서 목표값 T가 되는 쌍이 총 몇 개인지 찾는 것이다.

다음과 같은 방식으로 접근하였다.
1. 배열 A의 모든 부 배열의 합을 구하여 딕셔너리에 각 합이 몇 번 등장하는지 저장한다.
2. 배열 B의 모든 부 배열의 합을 구하면서, 각 합 s_b에 대해 T - s_b가 A의 합 딕셔너리에 있는지 확인한다.
3. 만약 T - s_b가 존재한다면, 그 값의 등장 횟수만큼 최종 결과에 더한다.

존재하지 않는 키에 접근하였을 때 0을 반환하기 위해 defaultdict를 사용하였다.