1/30 Coding Test

✅ Coding Test

🎈 2143 두 배열의 합

한 배열 A 내에서 부 배열은 A[i] ~ A[j]을 말한다. 이러한 부 배열의 합은 sum(A[i~j])를 의미한다. 각 원소가 정수인 두 배열 A, B가 주어졌을 때 A의 부 배열의 합에 B의 부 배열의 합을 더해 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램 만들기.

import sys
input = sys.stdin.readline
import bisect

t = int(input())
n = int(input())
arr_a = list(map(int, input().split()))
m = int(input())
arr_b = list(map(int, input().split()))

sum_a, sum_b = [], []
# 가능한 부분합 모두 구하기
for i in range(n):
    cnt = 0
    for j in range(i, n):
        cnt += arr_a[j]
        sum_a.append(cnt)
for i in range(m):
    cnt = 0
    for j in range(i, m):
        cnt += arr_b[j]
        sum_b.append(cnt)
sum_b.sort()
answer = 0

for i in sum_a:
    idx = t - i
    left = bisect.bisect_left(sum_b, idx)
    right = bisect.bisect_right(sum_b, idx)
    answer += right-left
print(answer)

< 해설 >

로직은 다음과 같다.
1. 주어진 두 배열 A, B의 모든 부분합을 구한다.
2. B 배열의 부분합을 구한 리스트를 정렬해 sum_a (a배열의 부분합)과 정렬된 b배열의 부분합을 이분 탐색해 합한 값이 t일 때의 경우의 수를 구한다.

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

이런 케이스의 경우 중복된 수가 존재할 때 bisect 라이브러리를 통해 파악해야 한다.
따라서 t-sum_a[i]값이 정렬된 sum_b에 존재한다면 t가 된다는 것이니 t-sum_a[i]값이 위치한 sum_b의 인덱스를 left, right 변수로 구한 후 경우의 수를 right - left로 더해준다.