누적합 - 2143번: 두 배열의 합

jisu_log·2025년 8월 7일

알고리즘 문제풀이

목록 보기
64/105


  • 배열의 접두합 P[k]는 A[0]부터 A[k−1]까지의 누적합을 저장해 구간합을 O(1)에 계산할 수 있게 함!
    -> 구간 A[i:j]의 합은 P[j] − P[i] 한 번의 뺄셈으로 바로 얻을 수 있음
  • A와 B의 모든 연속 부분합을 접두합으로 빠르게 뽑아내면 전체가 O(n² + m²) 시간에 가능해짐
  • collections.Counter로 각 부분합의 등장 횟수를 세어 두면, B에서 필요한 값(t−ai)을 바로 조회 가능
  • A 부분합 종류마다 “등장 횟수 × B에서의 대응 횟수”를 더해 합이 t인 쌍의 총개수를 효율적으로 구하기
import sys
from collections import Counter

input = sys.stdin.readline

t = int(input())
n = int(input())
A = list(map(int, input().split()))
m = int(input())
B = list(map(int, input().split()))

# 접두합 구하기
prefix_A = [0] * (n + 1)

for i in range(n):
    prefix_A[i + 1] = prefix_A[i] + A[i]


prefix_B = [0] * (m + 1)

for i in range(m):
    prefix_B[i + 1] = prefix_B[i] + B[i]


a_part_sum = []

for r in range(1, n + 1):
    for i in range(n):
        if i + r - 1 < n:
            a_part_sum.append(prefix_A[i + r] - prefix_A[i])

b_part_sum = []

for r in range(1, m + 1):
    for i in range(m):
        if i + r - 1 < m:
            b_part_sum.append(prefix_B[i + r] - prefix_B[i])


# Counter로 각 부분합 값이 몇 번씩 등장하는지 세기
a_cnt = Counter(a_part_sum)
b_cnt = Counter(b_part_sum)


cnt = 0

# A부분합 + B부분합 = t인 경우만 세어야 함
for ai, ca in a_cnt.items(): # ai: A의 부분합 중 하나의 값, ca: ai가 등장한 횟수
    cb = b_cnt[t - ai] # cb: (t - ai)가 B의 부분합에서 등장한 횟수
    cnt += ca * cb # 둘을 곱해서 가능한 모든 조합 갯수만큼 더해주기

print(cnt)

0개의 댓글