백준 2143 두 배열의 합 파이썬(python)

문제

한 배열 A [1], A [2], …, A [n]에 대해서, 부 배열은 A [i], A [i+1], …, A [j-1], A [j] (단, 1 ≤ i ≤ j ≤ n)을 말한다. 이러한 부 배열의 합은 A [i]+…+A [j]를 의미한다. 각 원소가 정수인 두 배열 A [1], …, A [n]과 B [1], …, B [m]이 주어졌을 때, A의 부 배열의 합에 B의 부 배열의 합을 더해서 T가 되는 모든 부 배열 쌍의 개수를 구하는 프로그램을 작성하시오.

예를 들어 A = {1, 3, 1, 2}, B = {1, 3, 2}, T=5인 경우, 부 배열 쌍의 개수는 다음의 7가지 경우가 있다.

T(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력

첫째 줄에 T(-1,000,000,000 ≤ T ≤ 1,000,000,000)가 주어진다. 다음 줄에는 n(1 ≤ n ≤ 1,000)이 주어지고, 그 다음 줄에 n개의 정수로 A [1], …, A [n]이 주어진다. 다음 줄에는 m(1 ≤ m ≤ 1,000)이 주어지고, 그다음 줄에 m개의 정수로 B [1], …, B [m]이 주어진다. 각각의 배열 원소는 절댓값이 1,000,000을 넘지 않는 정수이다.

출력

첫째 줄에 답을 출력한다. 가능한 경우가 한 가지도 없을 경우에는 0을 출력한다.

---

풀이

[알고리즘]

누적합을 이용한 뒤 이분 탐색을 이용해 풀 수 있었다. 다른 풀이를 찾아보니 굳이 이분탐색을 이용하여 풀 지 않고 누적합을 구하는 도중 개수를 k-v 값으로 저장해놓은 풀이도 존재했다. 후자로 푼 풀이가 시간 복잡도 면에서 더 효율적인 거 같다.  

 

반응형

 

[풀이 방법]

1. 입력 받은 an, bn 리스트 값들의 누적합을 구한 뒤 정렬 시켜준다.
2. 누적합 리스트를 각각 aan, bbn이라 한다.
3. 문제에서는 aan[i] + bbn [i] = T 값을 구해 달라고 했는데 즉, T - aan [i] = bbn [i] or T - bbn [i] = aan [i] 
4. 파이썬에서 제공하는 bisect 라이브러리를 이용하면 쉽게 리스트에 그 값이 존재하는지 index 값을 알 수가 있다.
5. tmp = [2, 3, 3, 3, 4, 4, 5, 6, 7] 이라는 리스트가 있을 때 bisect_left(tmp, 3), bisect_right(tmp, 3) 결과 값은 1, 4가 나올 것이다. 없는 값을 넣는다면 bisect_left(tmp, 8) 대략적인 index 위치인 9가 나오게 된다. 하지만 ans 값에 더할 때는 right - left 값을 더하기 때문에 없는 값을 조회한 left, right 값은 같기 때문에 0이 되고 더해도 의미가 없다.

import bisect
import sys

T = int(sys.stdin.readline())
a = int(sys.stdin.readline())
an = list(map(int, sys.stdin.readline().split()))
b = int(sys.stdin.readline())
bn = list(map(int, sys.stdin.readline().split()))

aan, bbn = [], []
for i in range(a):
    for j in range(i+1, a+1):
        aan.append(sum(an[i:j]))
for i in range(b):
    for j in range(i + 1, b + 1):
        bbn.append(sum(bn[i:j]))

aan.sort()
bbn.sort()

ans = 0
for i in range(len(aan)):
    bmp = T - aan[i]
    left = bisect.bisect_left(bbn, bmp)
    right = bisect.bisect_right(bbn, bmp)
    ans += (right - left)
print(ans)

---

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

2143번: 두 배열의 합