[백준 / 2143 / Python] 두 배열의 합

문제 요약

두 개의 정수 배열 A, B가 주어집니다.
각 배열에서 연속된 부분 구간(부 배열)을 하나씩 골라,
A의 부 배열의 합 + B의 부 배열의 합 = T가 되는 쌍의 개수를 구하는 문제입니다.

예를 들어,
A = [1, 3, 1, 2], B = [1, 3, 2], T = 5라면
A의 부 배열과 B의 부 배열을 각각 골라 합이 5가 되는 경우의 수를 모두 세야 합니다.

---

접근 방식

처음에는 모든 부 배열 쌍을 직접 만들어 합을 비교하는 완전탐색을 떠올릴 수 있습니다.
하지만 배열의 길이가 최대 1,000이므로,
O(N² × M²) 방식은 시간 초과가 발생합니다.

여기서 "부 배열의 합"이라는 조건에 주목하면,
각 배열의 모든 부 배열 합을 미리 구해놓고,
A의 부 배열 합 + B의 부 배열 합 = T
즉, A의 부 배열 합이 x라면, B의 부 배열 합이 T-x인 경우의 수를 세면 됩니다.

실생활 비유로 설명하면,
A에서 가능한 모든 합을 미리 적어두고,
B에서 가능한 합을 하나씩 꺼내어 "T를 만들려면 A에서 어떤 합이 필요하지?"를 찾아
그 개수를 곱해주는 방식입니다.

이때,

• A의 부 배열 합을 모두 구해 Counter(딕셔너리)로 빈도수를 저장
• B의 부 배열 합을 모두 구하며, T-해당합이 A에 몇 번 나오는지 누적
  이렇게 하면 효율적으로 정답을 구할 수 있습니다.

---

코드 구현

아래는 Python으로 구현한 코드입니다.

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()))

# 모든 부 배열의 합을 구하는 함수
def subarray_sums(arr):
  sums = []
  for i in range(len(arr)):
    total = 0
    for j in range(i, len(arr)):
      total += arr[j]
      sums.append(total)
  return sums

A_sums = subarray_sums(A)
B_sums = subarray_sums(B)

A_counter = Counter(A_sums)

result = 0
for b in B_sums:
  result += A_counter[T - b]

print(result)

---

시간복잡도 & 최적화

• 시간복잡도:
  • 각 배열의 부 배열 합 구하기: O(N²), O(M²)
  • Counter를 통한 조회 및 누적: O(N² + M²)
  • 전체적으로 O(N² + M²)로, N, M이 1,000일 때 충분히 통과 가능
• 최적화 포인트:
  • Counter(딕셔너리)를 사용해 A의 부 배열 합 빈도수를 미리 저장
  • B의 부 배열 합을 순회하며 T-합이 A에 몇 번 나오는지 바로 조회

---

느낀 점 / 팁

• "두 배열의 합" 유형은 한 쪽을 미리 모두 저장해두고,
  다른 쪽을 순회하며 보완값을 찾는 방식이 매우 효과적입니다.
• Counter(딕셔너리) 자료구조를 활용하면
  중복된 합의 빈도수까지 한 번에 처리할 수 있어
  효율성과 코드 가독성 모두 잡을 수 있습니다.
• 부 배열(연속 구간)의 합을 모두 구하는 패턴은
  다양한 누적합/투포인터 문제에서도 자주 등장하니 익혀두면 좋습니다.

---

참고

• 백준 2143번 문제 링크
• 파이썬 collections.Counter 공식 문서
• 누적합(Prefix Sum) 위키백과

---