[백준] 2143 두 배열의 합 풀이 Python

문제로 이동

---

---

접근 방법

누적합+이분탐색 카테고리여서 이분 탐색을 어떻게 사용할지 고민했습니다.

1. 전처리

배열 1개의 경우, 누적합을 구하기 위해 A[0] ~ A[n] 까지 배열을 순회하며 누적합을 넣어주지만 해당 문제를 이렇게 접근해버리면 안된다는 것은 알고 계실겁니다.

# 일반적인 배열 1개의 누적합
arr = [0,1,2,3,4,5,6,7,8,9]
for i in range(1,len(arr)):
	arr[i] = arr[i-1]+arr[i]
# result : [0,1,3,6,10,15,21,28,36,45]

2개의 배열이기 때문에 이런 의문이 생길 수 있습니다.

"각 배열의 비교 범위를 어떻게 선택할 것인가?"

B 배열의 탐색 index가 하나 증가하게 되면 A 배열의 index는 처음으로 돌아가 다시 탐색을 해야 하기 때문인데요,

이 것을 전처리 과정을 통해 해결할 수 있습니다.

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) = O(N^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)  = O(M^2)
    s = Barr[i]
    Bsum.append(s)
    for j in range(i+1,M):
        s+=Barr[j]
        Bsum.append(s)

시간 복잡도 : O(N²)

각 index에 대한 누적합을 미리 계산해주는 방식입니다. 모든 누적합이 미리 만들어져 있으니 하나의 Sum array를 기준으로 T 값을 맞추는 게 간편해졌죠

2. 이분 탐색

결국 우리는 T = Asum[i] + Bsum[j] 를 만족시키는 부분 배열을 찾아야 합니다.

다르게 생각하면 T - Asum[i] = Bsum[j]이 되는 것이죠.

 

이제, Bsum Array를 보고 이런 궁금증이 생깁니다.

Asum+Bsum=T에 해당되는 개수를 어떻게 구할 것인가?

여기서 사용되는 것이 이분 탐색입니다.

이분 탐색에 left와 right 또는 upper_bound, lower_bound를 들어보셨을 겁니다.

Python에서는 lower_bound가 bisect_left로, upper_bound가 bisect_right로 사용되니 표현을 bisect를 사용하겠습니다.

 

간단히 이 원리에 대해 설명하자면, 정렬된 [1,2,2,2,3]라는 Array에서

bisect_left(2)를 했다면 index 1을 반환하고

bisect_right(2)를 했다면 index 4를 반환하게 됩니다.

여기서 2의 개수는 right - left = 3과 동일합니다.

 

만약 조건문을 통해 하나씩 카운트했다면 O(N)만큼 소요되므로 O(logN)인 이분 탐색을 이용하는 거죠.

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

 

이 과정을 정렬된 Asum 개수만큼 연산하여 answer에 모든 값을 더해주게 됩니다.

---

Python 코드

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)