[누적합] 백준 2143번: 두 배열의 합
본문 바로가기

Coding Test

[누적합] 백준 2143번: 두 배열의 합

문제 : https://www.acmicpc.net/problem/2143
 
1. 모든 가능한 부분 배열의 합을 미리 계산하여 저장한다.
2. A의 부분 배열의 합과 B의 부분 배열의 합이 T가 되는 개수를 카운팅한다.
 
누적합 계산하는 이유
A = [1, 2, 3]
A_cumulate = [1, 3, 6, 2, 5, 3]
(1, 1+2, 1+2+3, 2, 2+3, 3)
이렇게 모든 가능한 부분 배열의 합을 미리 계산해두면, A와 B의 부분 배열의 합을 빠르게 찾을 수 있음 : O(1)
이 방식은 메모리를 더 사용하지만, 시간 복잡도를 최적화할 수 있는 전략
두 배열의 크기가 크지 않을 때 효과적인 해결 방법

import bisect

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

A_cumulate = []
B_cumulate = []

for i in range(n):  #O(A*(A-1)/2) = O(N^2)
    s = A[i]
    A_cumulate.append(s)
    for j in range(i+1,n):
        s += A[j]
        A_cumulate.append(s)

for i in range(m):  #O(B*(B-1)/2)  = O(M^2)
    s = B[i]
    B_cumulate.append(s)
    for j in range(i+1,m):
        s += B[j]
        B_cumulate.append(s)

A_cumulate.sort()
B_cumulate.sort()

ans = 0
for x in A_cumulate:
    l = bisect.bisect_left(B_cumulate, T-x)
    r = bisect.bisect_right(B_cumulate, T-x)
    ans += r-l
print(ans)