(Python) 백준 2143

백준 2143

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
2 초	64 MB	28157	9731	6524	31.708%

문제

한 배열 A[1], A[2], …, A[n]에 대해서, 부 배열은 A[i], A[i+1], …, A[j-1], A[j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A[i]+…+A[j]를 의미한다. 각 원소가 정수인 두 배열 A[1], …, A[n]과 B[1], …, B[m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2]

입력

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A[1], …, A[n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그 다음 줄에 m개의 정수로 B[1], …, B[m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

---

답안

import sys
from collections import defaultdict

target = int(sys.stdin.readline())
_ = int(sys.stdin.readline())
list_a = list(map(int, sys.stdin.readline().split()))
_ = int(sys.stdin.readline())
list_b = list(map(int, sys.stdin.readline().split()))
counts = defaultdict(lambda: 0)
result = 0

for i in range(len(list_a)):
    local_sum = 0
    for j in range(i, len(list_a)):
        local_sum += list_a[j]
        counts[local_sum] += 1

for i in range(len(list_b)):
    local_sum = 0
    for j in range(i, len(list_b)):
        local_sum += list_b[j]
        if counts[target - local_sum] > 0:
            result += counts[target - local_sum]

print(result)

풀이

단순 구현을 해서 풀었다.

이렇게 문제가 쉽게 풀릴까 싶어 인터넷의 다른 풀이를 찾아봤었는데, 이분탐색을 이용한 풀이가 대부분이었다.

내가 풀었던 방식을 설명해보면,

1. 부배열 A에 대해 $O(n^2)$으로 탐색하며 { key(부분수열의 합): value(해당 부분수열의 합이 만들어지는 경우의 수) } 형식으로 데이터를 만든다.
2. 부배열 B에 대해 위와 같은 방식으로 탐색하며, 부배열 A의 부분수열의 합이 target-부배열 B의 부분수열의 합과 같아지는 경우, 그 경우의 수를 정답에 더해준다.

문제에서 주어진 시간이 2초였기 때문에, 길이 최대 1000짜리 리스트를 2번 $O(n^2)$ 탐색하는 것은 큰 문제 없을 것이라고 생각했고, 그렇게 풀었더니 정답 판정을 받았다.

데이터의 크기가 1000으로 크지 않아 이런 풀이도 가능한 것이었지만, 데이터의 크기가 커진다면 이분탐색을 이용해 더 빠르게 문제의 답을 구할 수 있을 것으로 보인다.