boj.kr/2143 두 배열의 합 (gold3) 파이썬 풀이

1. 첫번째 배열에서 가능한 모든 부분배열의 합을 N_sums 리스트에 저장한다.

2. 두번째 배열에서 가능한 모든 부분배열의 합을 M_sums 리스트에 저장한다.

3. 첫번째 배열의 모든 요소를 순회하면서 아래 3.1.-3.3 을 반복한다.

3.1. M_sums에서 bisect_right으로 T-첫번째 배열의 요소 값을 탐색한 인덱스를 r에 저장한다.

3.2. M_sums에서 bisect_left으로 T-첫번째 배열의 요소 값을 탐색한 인덱스를 l에 저장한다.

3.3. r-l을 cnt변수에 증감한다.

4. 모든 가능한 경우의 수 cnt를 출력한다.

 

T = int(input())
N = int(input())
N_list = list(map(int, input().split()))
M = int(input())
M_list = list(map(int, input().split()))

N_sums = []
for i in range(len(N_list)):
    for j in range(i+1, len(N_list)+1):
        N_sub = N_list[i:j]
        N_sums.append(sum(N_sub))

M_sums = []
for k in range(len(M_list)):
    for l in range(k+1, len(M_list)+1):
        M_sub = M_list[k:l]
        M_sums.append(sum(M_sub))

M_sums.sort()

from bisect import bisect_left, bisect_right

cnt = 0
for N_sum in N_sums:
    l = bisect_left(M_sums, T-N_sum)
    r = bisect_right(M_sums, T-N_sum)
    cnt += r-l

print(cnt)

 

첫번째 배열과 두번째 배열의 모든 부분배열의 합들의 목록을 각각 구해야 한다. 그러나 두 부분배열의 합 목록에 대해 다시 이중 for문을 돌리지 않아도 된다. 왜냐하면 두번째 배열의 모든 부분배열의 합 목록에서 T-첫번째 배열의 요소를 만족하는 값이 몇개인지를 이분탐색으로 구할 수 있기 때문이다.

 

bisect_right는 요소가 있는 경우 해당 인덱스+1을 반환하고 요소가 없는 경우 오름차순으로 들어갈 위치를 반환한다.

bisect_left는 요소가 있는 경우 해당 인덱스를 반환하고 요소가 없는 경우 오름차순으로 들어갈 위치를 반환한다.

 

그래서 정렬된 리스트에 대해서 bisect_right값 - bisect_left값은 찾으려는 요소가 없으면 0이고 찾으려는 요소가 있으면 찾으려는 요소의 총 개수이다.