백준 2143 : 두 배열의 합 (Python)

본문 링크

import sys
input=sys.stdin.readline
import bisect
T=int(input())
N=int(input())
Prefix_sum_1=list(map(int,input().split()))
M=int(input())
Prefix_sum_2=list(map(int,input().split()))

dp1=[] ; dp2=[]

for i in range(N):
    tmp=0
    for j in range(i,N):
        tmp+=Prefix_sum_1[j]
        dp1.append(tmp)

for i in range(M):
    tmp=0
    for j in range(i,M):
        tmp+=Prefix_sum_2[j]
        dp2.append(tmp)
dp2.sort() ; result=0

for i in dp1:
    temp=T-i
    left=bisect.bisect_left(dp2,temp)
    right=bisect.bisect_right(dp2,temp)
    result+=right-left
print(result)
"""
dp2에 temp 값이 존재한다면 경우의 수가 하나 늘어난다. 여러개가 있을수 있기 때문에 bisect를 사용해서  
모든 경우를 찾아준다.
"""

📌 어떻게 접근할 것인가?

오래 생각해봤는데 아무래도 가장 깔끔한 방법은
$dp1$ 배열과 $dp2$ 배열에 각각 모든 연속 부분수열을 넣는것이다.
이후 $dp2$ 배열을 정렬해서 for i in range(len(dp1)) 과 정렬된 $dp2$ 값을 이분탐색해서 두 값의 합이 T 일때 경우의 수를 찾는것이다.

하지만 아주 까다로운 반례가 있었다.

• 입력
  2
  1
  1
  10
  1 1 1 1 1 1 1 1 1 1
• 답
  10

즉 , 중복된 수가 존재할수 있다는것이다.

따라서 우리는 bisect 라이브러리를 통해서
T-dp1[i] 값이 $dp2$ 값에 존재한다면 dp1[i]+dp2[어딘가에있을 값]의 합이 T 가 가능하다는 것이니 T-dp1[i] 값이 위치한 dp2 인덱스를 bisect_left 와 bisect_right 를 구해서
중복된 값을 함께 계산해준다.

이렇게 접근하면 중복된 수를 처리 가능하다.

✅ 코드에서 주의해야할 부분

• dp2만 정렬한다.