[코딩 테스트 연습(파이썬/Python)] 백준 2143번 _ 두 배열의 합

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

2143번: 두 배열의 합

[문제]

한 배열 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) a_에 저장된 모든 누적 합(x)에 대하여, b_에 t-x가 존재한다면 그 갯수를 정답에 더해준다.

3) 2)를 빠르게 하기 위해 Counter 모듈로 b_의 원소의 갯수를 센다.

 

from collections import Counter
import sys
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()))

a_ = []; b_ = []
for i in range(n):
    x = a[i]
    a_.append(x)
    for j in range(i+1,n):
        x += a[j]
        a_.append(x)
for i in range(m):
    y = b[i]
    b_.append(y)
    for j in range(i+1,m):
        y += b[j]
        b_.append(y)
b_ = Counter(b_)        
ans = 0
for x in a_:
    if t-x in b_:
        ans += b_[t-x]
print(ans)

이분 탐색보다 빠른 풀이인 것 같다. 원트 뿌듯 :)