[BOJ, 백준] 2143 두 배열의 합 python

 

 

 

2개의 누적합을 얼마나 잘 쓸 수 있겠니? 라는 문제인데

 

사실 누적합과 딕셔너리만 있어도 풀 수 있는 문제 같다

 

(키워드에는 이분 탐색이 있길래...)

 

물론, 이분 탐색을 통해서 빠르게 값을 찾아내는 것도 중요하겠다.

 

대신 이 경우에는 딕셔너리가 아닌 값을 사용해야겠지만 ...

 

답안 코드는 아래와 같다

 

from collections import defaultdict
T = int(input())
N = int(input())
A = list(map(int,input().split()))
M = int(input())
B = list(map(int,input().split()))


psum_A = [0 for _ in range (0,N+1)]
psum_B = [0 for _ in range (0,M+1)]
for i in range (1,N+1):
    psum_A[i] = psum_A[i-1] + A[i-1]
for i in range (1,M+1):
    psum_B[i] = psum_B[i-1] + B[i-1]

partial_sum_A = defaultdict(int)
partial_sum_B = defaultdict(int)

for i in range (1, N+1):
    for j in range (0,i):
        partial_sum_A[psum_A[i] - psum_A[j]] += 1

for i in range (1, M+1):
    for j in range (0,i):
        partial_sum_B[psum_B[i] - psum_B[j]] += 1

pnum_A = list(partial_sum_A.keys())
pnum_A.sort()
cnt = 0

for i in range (0,len(pnum_A)):
    target = T - pnum_A[i]
    if target in partial_sum_B:
        cnt += partial_sum_A[pnum_A[i]] * partial_sum_B[target]
print(cnt)