[백준 파이썬] 2143 두 배열의 합

백준 2143번

메모리 제한이 64MB라 딕셔너리를 사용해서 풀면 메모리 초과가 발생할 줄 알았지만 의외로 통과되었다. 원래 의도는 이분 탐색을 이용해서 푸는 문제다.

누적 합 배열을 구하고 그 누적 합 배열에서 값 두개를 선택하는 것으로 누적 합의 범위를 지정한다고 볼 수 있다. 다음으로 하나의 누적 합 배열의 원소 개수를 카운트 하고 타겟 값 T에서 특정 범위 합을 뺀 값이 이 카운트 딕셔너리에 몇 개나 있는지 세면 된다.

from collections import Counter

T = int(input())
N = int(input())
A = list(map(int, input().split()))
M = int(input())
B = list(map(int, input().split()))

# A 배열의 전체 누적 합 배열
suf_A = [0 for _ in range(N+1)]
for i in range(1, N+1):
    suf_A[i] = suf_A[i-1] + A[i-1]

# B 배열의 전체 누적 합 배열
suf_B = [0 for _ in range(M+1)]
for i in range(1, M+1):
    suf_B[i] = suf_B[i-1] + B[i-1]

# A 배열의 모든 가능한 범위 합
part_A = []
for i in range(N+1):
    for j in range(i+1, N+1):
        part_A.append(suf_A[j] - suf_A[i])
part_A.sort()

# B 배열의 모든 가능한 범위 합
part_B = []
for i in range(M+1):
    for j in range(i+1, M+1):
        part_B.append(suf_B[j] - suf_B[i])
part_B.sort()

# 메모리 사용 제한으로 유일한 값이 적은 배열로 카운트 딕셔너리 생성
if set(part_B) > set(part_A):
    part_A, part_B = part_B, part_A
count_B = Counter(part_B)

# T에서 A 범위 합을 뺀 값이 B 범위 합 배열에 몇 개가 있는지 카운트
ans = 0
for n in part_A:
    ans += count_B[T-n]

print(ans)