[백준 2143] 두 배열의 합 (python)

2143번: 두 배열의 합

---

 

S[j] - S[i-1] 을 통해 i부터 j까지의 합을 구할 수 있다. A의 합을 dictionary에 넣고 B와 T가 되는 A의 합을 찾자!

 

sumA에 A배열의 합, sumB에 B배열의 합을 저장한다.

Adict에 A의 부분합을 저장한다. key=부분합 value=개수

B의 부분합을 구하면서 그 부분합과 T를 만들 수 있는 A의 합을 Adict에서 찾는다.

 

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

Adict = defaultdict(int)
sumA = [0]
sumB = [0]
for a in A:
    sumA.append(sumA[-1]+a)
for b in B:
    sumB.append(sumB[-1]+b)

for i in range(1, len(A)+1):
    for j in range(i):
        Adict[sumA[i]-sumA[j]] += 1

result = 0
for i in range(1, len(B)+1):
    for j in range(i):
        result += Adict.get(T-(sumB[i]-sumB[j]), 0)

print(result)