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

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

2143번: 두 배열의 합

 

import sys
input = sys.stdin.readline

if __name__ == '__main__':
    T = int(input())
    n = int(input())
    A = list(map(int, input().split()))
    m = int(input())
    B = list(map(int, input().split()))

    # Set sums
    a_sums = {}
    for i in range(n):
        cur_sum = 0
        for j in range(i, n):
            cur_sum += A[j]

            if a_sums.get(cur_sum):
                a_sums[cur_sum] += 1
            else:
                a_sums[cur_sum] = 1

    b_sums = {}
    for i in range(m):
        cur_sum = 0
        for j in range(i, m):
            cur_sum += B[j]

            if b_sums.get(cur_sum):
                b_sums[cur_sum] += 1
            else:
                b_sums[cur_sum] = 1

    # Solution
    count = 0
    for k, v in a_sums.items():
        if b_sums.get(T-k):
            count += v*b_sums[T-k]

    print(count)

 

배열의 크기가 1000이 최대이므로 부 배열의 합을 모두 저장할 수 있다. 예를 들어, 1 3 1 2의 경우 부배열의 합은 다음과 같이 된다.

이것을 dictionary에 저장한다. 위 경우를 딕셔너리에 저장한 결과는 다음과 같다.

{1: 2, 4: 2, 5: 1, 7: 1, 3: 2, 6: 1, 2: 1}

 

이제 원하는 값을 쉽게 구할 수 있다. 가령, 위 딕셔너리에서 1에 대응하는 부분 합의 경우의 수를 구하고 싶다고 하자. 그렇다면, 또 다른 딕셔너리에서 T - 1(=4)의 갯수를 가져온 뒤, 여기에 2를 곱해주면 된다!