문제

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

요약

두 배열 A, B가 주어졌을 때, 연속하는 부분 배열의 합을 더해서 T가 되는 모든 연속하는 부분 배열 쌍의 개수를 구하는 프로그램을 작성하라.

분류

  • 누적합
  • 해시 맵

풀이

1. 내 풀이 (누적합, 해시 맵)

N의 크기가 크기 때문에 O(NlogN)에 해결해야 한다. 또, 음수가 들어갈 수 있기 때문에 초깃값 설정에 유의하자.


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

prefix_sum_a = [0]
prefix_sum_b = [0]

d_a = dict()
d_b = dict()

acc_a = 0
for i in range(n):
    acc_a += arr_a[i]
    for val in prefix_sum_a:
        try:
            d_a[acc_a - val] += 1
        except:
            d_a[acc_a - val] = 1
    prefix_sum_a.append(acc_a)

acc_b = 0
for i in range(m):
    acc_b += arr_b[i]
    for val in prefix_sum_b:
        try:
            d_b[acc_b - val] += 1
        except:
            d_b[acc_b - val] = 1
    prefix_sum_b.append(acc_b)

ans = 0
l = list(d_a.keys())
for key in l:
    try:
        ans += d_a[key] * d_b[t - key]
    except:
        continue

print(ans)