[백준 2143번] 파이썬 - 두 배열의 합

http://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을 넘지 않는 정수이다.

 

# 접근 방법

• 누적합과 이분 탐색을 이용해주면 될 것 같다.
• 2개의 배열을 탐색해야 하므로 각 index에 대한 누적합을 미리 계산해준다 -> O(N^2)
  • 즉, 모든 누적합이 미리 만들어져 있으니 하나의 Sum array를 기준으로 T 값을 맞춰준다.
• 이분 탐색의 경우 Python에서는 lower_bound가 bisect_left로, upper_bound가 bisect_right로 사용되므로 bisect을 이용해준다. -> O(logN)

import bisect  
T = int(input())  
N = int(input())  
Aarr = list(map(int,input().split()))  
M = int(input())  
Barr = list(map(int,input().split()))  
Asum = []  
Bsum = []  
for i in range(N):  #O(A*(A-1)/2)  
    s = Aarr[i]  
    Asum.append(s)  
    for j in range(i+1,N):  
        s+=Aarr[j]  
        Asum.append(s)  
for i in range(M):  #O(B*(B-1)/2)  
    s = Barr[i]  
    Bsum.append(s)  
    for j in range(i+1,M):  
        s+=Barr[j]  
        Bsum.append(s)  
Asum.sort()  
Bsum.sort()  
answer = 0  
for i in range(len(Asum)):  
    l = bisect.bisect_left(Bsum,T-Asum[i])  
    r = bisect.bisect_right(Bsum,T-Asum[i])  
    answer+=r-l  
print(answer)