[🥇3 / 백준 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(-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을 출력한다.

 

예제 입력  5 4 1 3 1 2 3 1 3 2

예제 출력  7

 

풀이

 

딕셔너리 활용, 이분 탐색, 브루트 포스, 누적 합 등 다양한 알고리즘으로 해결이 가능한 문제이다.

그 중에서 브루트 포스와 이분 탐색을 이용해 문제를 해결했다.

 

우선 브루트 포스로 A배열과 B배열의 모든 가능한 조합의 합을 구한다.

각각 배열의 크기가 1,000에 해당하니 O(N^2) 이라 해도 최대 2,000,000 연산 안에 구할 수 있다.

 

그리고 B를 정렬을 정렬한 뒤, A를 순환하며 B에서 이분 탐색을 진행하며 합이 T가 되는 경우를 찾는다.

 

import sys, bisect
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()))

# [sum(a[i:j]) for i in range(n) for j in range(i+1,n+1)] 와 같이 구할 수도 있지만,
# sum을 매번 반복해야 하기 때문에 시간이 오래 걸리는 것 같다.
# 때문에 브루트 포스를 하며 누적 합 방식으로 모든 조합의 합을 구한다.
A = []
for i in range(n):
    tmp = 0
    for j in range(i,n):
        tmp += a[j]
        A.append(tmp)

B = []
for i in range(m):
    tmp = 0
    for j in range(i,m):
        tmp += b[j]
        B.append(tmp)
B.sort()

cnt = 0
for i in A:
    diff = t-i
    idx = bisect.bisect_left(B,diff)
    if idx >= len(B):
        continue
    
    # 조합의 합이 같은 경우가 있을 수 있기 때문에 그 인덱스의 차이를 구해준다.
    idx_right = bisect.bisect_right(B,diff)
    if B[idx] == diff:
        cnt+=idx_right-idx
print(cnt)