BOJ #2143 두 배열의 합

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

문제 2143. 두 배열의 합

문제 풀이

처음엔 부분합 문제로 A와 B의 모든 부분 배열의 합을 배열로 저장하고 A의 모든 원소에 대해 탐색 배열을 B의 부분배열의 합 배열로 두고 타겟 - A의 부분합 배열의 원소값을 이진 탐색으로 찾아서 해결할 수 있을 줄 알았다. 그야 메모리 제한이 64MB라서 공간복잡도가 작아보였기 때문...

그러나 이 문제는 A의 부분합을 배열로 생성하고 B의 부분합은 dictionary로 생성하여 탐색 과정을 $O(1)$로 줄여야 하는 문제였다. C계열 언어였다면 통과했을 지 몰라도 python에서는 시간 초과가 발생하기 때문이다. 그럴만도 한게 부분합 배열은 엄연히 $O(n^2)$이고 매번 B의 배열을 이진탐색으로 찾으면 탐색과정은 대충 $2log\ n$이 되는데 상당히 버거울 만도 하다. A의 부분합 배열에 대해 순차적으로 탐색하여 B 부분합 딕셔너리에 타겟 - A 부분합 배열의 원소 키를 조회하여 더하면 되겠다.

python 코드

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_A, sum_B = [], defaultdict(int)
for i in range(1, N):
    l, r = 0, i
    buf = sum(A[l:r])
    sum_a.append(buf)
    while r < N:
        buf -= A[l]
        buf += A[r]
        sum_a.append(buf)
        l += 1
        r += 1
for i in range(1, M):
    l, r = 0, i
    buf = sum(B[l:r])
    sum_b[buf] += 1
    while r < M:
        buf -= B[l]
        buf += B[r]
        sum_b[buf] += 1
        l += 1
        r += 1
sol = 0
for s_a in sum_a:
    sol += sum_b[T - s_a]
print(sol)