[ 백준 ][ 골드3 ] 2143번 - 두 배열의 합 ( 파이썬 PYTHON )

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

2143번: 두 배열의 합

 

코드

---

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()))

sum_A, sum_B = [0], [0]
for i in range(n):
    sum_A.append(sum_A[i] + A[i])
for i in range(m):
    sum_B.append(sum_B[i] + B[i])

dic = defaultdict(int)
for i in range(n + 1):
    for j in range(i + 1, n + 1):
        dic[sum_A[j] - sum_A[i]] += 1

result = 0
for i in range(m + 1):
    for j in range(i + 1, m + 1):
        result += dic[T - (sum_B[j] - sum_B[i])]

print(result)

 

풀이

---

기본적인 부분합 문제였다면 부분 합이 T인 경우의 수를 구하는 것이었겠지만 이 문제는 좀 더 어렵게 꼬아서 두 배열의 부분 합의 합이 T인 경우의 수를 구해야 한다.

여기서 경우의 갯수만 출력하면 되므로 모든 부분합의 원소들을 저장할 필요는 없다.

 

< 모든 부분합 탐색 - 시간초과 >

처음에는 A,B 두 집합의 모든 부분합을 검사하면서 T이하인 부분합의 합을 배열에 모두 저장하고 그렇게 만들어진 두 배열의 탐색하면서 합이 T이면 정답을 1씩 올려주었다.

이렇게 하면 모든 경우의 수를 다 탐색하고 저장한 것이므로 당연하게도 시간초과가 발생했다.

 

< 딕셔너리 사용 - 통과 >

먼저 A와 B의 누적합 배열을 만든다. - sum_A, sum_B

A의 모든 부분합을 탐색하면서 딕셔너리에 부분합을 인덱스로 하여 해당 값이 나올 때마다 1씩 더해준다.

 

이 때, 딕셔너리를 이용하는 이유는 딕셔너리의 값에 접근할 때는 O(1) 시간복잡도가 필요하고 defaultdict 라이브러리를 사용하면 딕셔너리 키값이 없어도 자동으로 추가해주기 때문이다. - dic = defaultdict(int)

 

B의 누적합을 모두 탐색하며 위에서 만들어 놓은 부분합과 더해서 T가 되는 경우를 result에 더해준다.

이 때, T - B의 부분합 = A의 부분합이므로 dic[T-B의 부분합]의 값을 더해주면 된다.

defaultdict 라이브러리를 사용했기 때문에 만약 A의 부분합이 없다면, 즉 dic의 키 값이 없다면 0이 더해진다.