💡 문제
한 배열 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(-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을 출력한다.
💭 접근
이 문제는 두 배열의 부분합을 더하여 T가 되는 경우의 수를 구하는 문제로, 누적합을 구하여 쉽게 해결할 수 있다. 이때 각 배열의 원소는 자연수가 아닌 정수이므로 투 포인터 알고리즘은 사용할 수 없다.
먼저 배열 A와 배열 B에 대해서 누적합을 구한다.
p_sum_a = [0]
tmp = 0
for i in range(n):
tmp += a[i]
p_sum_a.append(tmp)
p_sum_b = [0]
tmp = 0
for i in range(m):
tmp += b[i]
p_sum_b.append(tmp)이후, 딕셔너리를 사용하여 먼저 배열 A의 누적합을 활용하여 구할 수 있는 모든 부분합의 경우의 수를 더해준다.
cnt = defaultdict(int)
for i in range(n):
for j in range(i + 1, n + 1):
cnt[p_sum_a[j] - p_sum_a[i]] += 1이제는 배열 B의 부분합을 사용할 차례인데, 이때는 목표 숫자인 T에서 배열 B의 부분합을 뺀 값의 경우의 수를 정답에 더해주면 된다.
ans = 0
for i in range(m):
for j in range(i + 1, m + 1):
ans += cnt[t - (p_sum_b[j] - p_sum_b[i])]위와 같이 구하면, 배열 A의 부분합과 배열 B의 부분합을 더한 수가 T가 될 때의 경우의 수를 구할 수 있다.
📒 코드
from collections import defaultdict
t = int(input())
n = int(input())
a = list(map(int, input().split()))
m = int(input())
b = list(map(int, input().split()))
p_sum_a = [0]
tmp = 0
for i in range(n):
tmp += a[i]
p_sum_a.append(tmp)
p_sum_b = [0]
tmp = 0
for i in range(m):
tmp += b[i]
p_sum_b.append(tmp)
cnt = defaultdict(int)
for i in range(n):
for j in range(i + 1, n + 1):
cnt[p_sum_a[j] - p_sum_a[i]] += 1
ans = 0
for i in range(m):
for j in range(i + 1, m + 1):
ans += cnt[t - (p_sum_b[j] - p_sum_b[i])]
print(ans)💭 후기
누적합을 활용하여 부분합을 구한 뒤, 딕셔너리에 경우의 수를 저장하는 재미있는 문제였다.
🔗 문제 출처
https://www.acmicpc.net/problem/2143
