[ BOJ ] 2143 : 두 배열의 합 ( GOLD 3 ) / Python

문제

한 배열 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을 출력한다.

풀이 과정

1. A 배열과 B 배열의 누적 합 배열을 만든다.

2. B 누적 합 배열을 전부 검토하면서 가능한 모든 연속 구간의 누적합을 얻는다. 똑같은 합을 가진 중복된 구간을 딕셔너리로 저장한다. ( 탐색 시 정확하게 찾기 위해 )

3. B 배열의 합 정보를 저장한 딕셔너리를 합 기준으로 정렬한다. 

4. A 누적 합 배열의 가능한 모든 연속 구간을 검토하면서, T와의 차이를 계산해 이 차이와 정확히 들어맞는 B 합이 있는지 이분 탐색한다.

5. 가능한 쌍의 개수를 카운트한다.

import sys
input = sys.stdin.readline

t = int(input())
n = int(input())
arrA = list(map(int, input().split()))
m = int(input())
arrB = list(map(int, input().split()))

arrASum = []
arrBSum = []

for i in range(n): 
    if i == 0: arrASum.append(arrA[i])
    else: arrASum.append(arrASum[-1] + arrA[i])
    
for i in range(m): 
    if i == 0: arrBSum.append(arrB[i])
    else: arrBSum.append(arrBSum[-1] + arrB[i])

arrB_dict = {}
for i in range(len(arrBSum)):
    if arrB_dict.get(arrBSum[i], -1) == -1: arrB_dict[arrBSum[i]] = 1
    else: arrB_dict[arrBSum[i]] += 1
    
    if i == 0: continue
    
    for j in range(i):
        if arrB_dict.get(arrBSum[i] - arrBSum[j], -1) == -1: arrB_dict[arrBSum[i]-arrBSum[j]] = 1
        else: arrB_dict[arrBSum[i] - arrBSum[j]] += 1

arrBoc = list(arrB_dict.items())
arrBoc.sort(key=lambda x:x[0])

def binary_search(k):
    global arrBoc

    start = 0
    end = len(arrBoc) - 1
    while start <= end:
        mid = (start + end) // 2

        if arrBoc[mid][0] == k: return arrBoc[mid][1]
        elif arrBoc[mid][0] < k: start = mid + 1
        else: end = mid - 1
    
    return 0

result = 0
for i in range(len(arrASum)):
    key = t - (arrASum[i])
    result += binary_search(key)
    if i == 0: continue
    
    for j in range(i):
        key = t - (arrASum[i] - arrASum[j])
        result += binary_search(key)

print(result)