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

[Gold III]

 

https://www.acmicpc.net/problem/2143

2143번: 두 배열의 합

 

풀이

부 배열의 합.. -> Prefix Sum!

우선 prefix sum을 구현하기 위해.. S_k (1~k번째 합)를 저장할 배열을 각각 생성.

 

그 다음이 문제인데...

우선 배열의 원소 조건이 절댓값이 1,000,000을 넘지 않는 정수 이기 때문에.. A[k+1] >= A[k] 가 아니다.

따라서 값이 들쭉날쭉할 수 있기 때문에.. pointer를 이용하는 느낌으로 가는 방향은 불가능.

 

고민을 좀 하다.. 모든 prefix sum에 대해 생각하면 TLE가 날까.. 생각해봤다.

n, m 모두 1,000 이하이기 때문에.. 한 배열에 대한 모든 prefix sum의 개수는 ₁₀₀₀C₂ = 약 500,000개정도.

A와 B의 모든 prefix sum에 대해 Brute-Force 식으로 생각한다면.. 당연히 50만x50만=2500억개.. 니까 TLE가 나겠지만,

우리는, 가능한 '경우의 개수' 를 생각하는 거니까, dictionary에 각각의 부분합의 결과에 따라 '개수'를 센다면,

50만+50만 = 100만개 정도로 끝낼 수 있지 않을까? 라는 아이디어.

---

예를 들자면, A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5일 때,

A의 부분합은 [1, 4, 5, 7], B의 부분합은 [1, 4, 6]이므로,

A의 prefix sum 들은 [1, 4, 5, 7 / 3, 4, 6 / 1, 3 / 2] 총 10개가 있으므로,

{1: 2, 2: 1, 3: 2, 4: 2, 5: 1, 6: 1, 7: 1} 과 같이 저장. ( {합: 개수} )

B의 prefix sum 들은 [1, 4, 6 / 3, 5 / 2] 총 6개이므로, {1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1} 과 같이 저장.

 

이렇게 하고, {1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1} 에서 key값으로 순회를 돌며, 계산.

예를 들면 T = 5이니까, key = 2일 때는 B에서의 prefix sum이 2이므로, A에서의 prefix sum은 5 - 2 = 3이 되어야 함.

B의 prefix sum이 2인 경우는 1가지, A의 prefix sum이 3인 경우는 2가지이므로, 경우의 수는 총 1 x 2 = 2가지.

이런 식으로 한 dictionary의 모든 key들에 대해 반복하면 된다.

 

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

# 부분합
As = [0]
for i in range(n):
    As.append(As[i] + A[i])

Bs = [0]
for i in range(m):
    Bs.append(Bs[i] + B[i])

# Prefix Sum
PA = {}
PB = {}

for i in range(n+1):
    for j in range(i):
        k = As[i] - As[j]
        if PA.get(k, -1) == -1:
            PA[k] = 1
        else:
            PA[k] += 1

for i in range(m+1):
    for j in range(i):
        k = Bs[i] - Bs[j]
        if PB.get(k, -1) == -1:
            PB[k] = 1
        else:
            PB[k] += 1

result = 0
# find by all dictionary keys
for i in PA:
    n = T - i
    if PB.get(n, -1) != -1:
        result += PA[i]*PB[n]

print(result)