파이썬) 두 배열의 합 - BOJ

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

2143번: 두 배열의 합

 

 

문제가 원하는바는 다음과 같다.

첫번째 배열 arr1 i부터 ,j까지의 합과 arr2의 k,l까지의 합이 t인 개수를 구하는 문제이다.

이 문제는 어떻게 풀어나가면 좋을까

 

단순히 생각해보자

arr1의 0번째부터. 0번째, 0번째부터 1번째 ... 1부터 1까지, 1부터 2까지 .... n부터 n까지

arr2의 0 - 0 ..... m - m까지 모든 경우의 수를 구하기 위해선 n * n * m *m이라는 어마어마한 수가 필요하다.

이 시간복잡도를 줄이는 것이 관건이다.

 

우선 i - j번째를 누적합으로 계산해보자

 

dic = defaultdict(int)

for i in range(n):
    for j in range(i, n):
        num = sum(arr1[i:j+1])
        dic[num] += 1

 

이중 반복문을 돌며 배열의 누적합을 계산해 딕셔너리에 저장한다.

이를 어따 쓰냐고 생각하겠지만 arr[i:j+1]의 합이

arr2의 누적합을 t에서 뺀 값과 같은지 비교만하면 끝이난다..

 

from collections import defaultdict

t = int(input())
n = int(input())

arr1 = list(map(int, input().split()))

m = int(input())

arr2 = list(map(int, input().split()))

dic = defaultdict(int)

for i in range(n):
    for j in range(i, n):
        num = sum(arr1[i:j+1])
        dic[num] += 1

answer = 0

for i in range(m):
    for j in range(i, m):
        num = t - sum(arr2[i:j+1])
        answer += dic[num]

print(answer)