[백준 2143] 두 배열의 합 (Gold 3)

문제

https://www.acmicpc.net/problem/2143

접근

1. 완전 탐색을 이용해 각각 A와 B에서 나올수 있는 모든 합을 오름차순으로 리스트에 따로 저장한다. (아래의 코드에는 sumA, sumB)
   1-1. 예를 들어, A= [1 3 1 2]일 때 sumA에 들어있는 값은 [1, 1, 2, 3, 3, 4, 4, 5, 6, 7]이다.
2. A의 부 배열의 합을 numA, b의 부 배열의 합을 numB라 했을 때, T = numA + numB라는 식이 성립된다. 즉 T - numA = numB 라는 식을 이용할 수 있다.
3. sumB에서 T - numA의 값으로 이분 탐색을 한다.
4. sumB에서 값이 중복되는 경우도 있다. 이를 찾기 위해 lower bound, upper bound의 성질을 이용한다.
   4-1 . [1,2,2,2,5] 라는 배열에서 2가 몇 개 있는지 찾는다고 해보자. 이 배열에서의 lower bound는 1, upper bound는 4가 된다. 즉 upper bound - lower bound는 2의 개수가 된다.

Code

import sys
from bisect import bisect_left, bisect_right


def generate(arr, n):
    ret = []
    for i in range(n):
        ret.append(arr[i])
        tmp = arr[i]
        for j in range(i + 1, n):
            tmp += arr[j]
            ret.append(tmp)

    return sorted(ret)


t = int(sys.stdin.readline())
a, arrA = int(sys.stdin.readline()), list(map(int, sys.stdin.readline().split()))
b, arrB = int(sys.stdin.readline()), list(map(int, sys.stdin.readline().split()))

sumA, sumB = generate(arrA, a), generate(arrB, b)
ans = 0

for numA in sumA:
    right = bisect_right(sumB, t - numA)
    left = bisect_left(sumB, t - numA)
    ans += right - left
print(ans)

여담

• 모듈이 편하긴 편하다.