[BOJ] 2143번 : 두 배열의 합(구간 합, 투포인터)

접근

부 배열은 수열의 연속된 부분에 대한 합이다.

다른 두개의 수열 각각의 부 배열에 대한 합이 특정한 값과 같은 조합의 갯수를 세는 문제다. 

 

두 수열에 대해 가능한 부 배열의 합을 모두 구하고, 투포인터를 이용해 각각의 합에 대해서 특정한 값과 맞는지 갯수를 세는 방식으로 문제를 해결했다.

 

해결

부 배열의 합을 구하기 위해서, 구간합을 먼저 구한 뒤 이중반복문으로 존재하는 모든 부 배열의 합을 구했다.

나는 부 배열의 합을 dict형식으로 저장해, '부 배열의 합'을 Key로, '갯수'를 Value로 사용하여 한번에 처리하도록 구성했다.

 

그리고 한 부 배열 조합은 부 배열의 합의 종류에 대해서 오름차순으로, 다른 부 배열 조합은 내림차순으로 정렬한 뒤에 투포인터를 이용해 각자의 인덱스로 접근하여 합의 탐색을 진행했다.

 

문제의 분류가 이분탐색으로 되어 있던데, 크기 순으로 정렬한다면 이분탐색도 당연히 사용해볼 수 있다.

 

코드

from collections import defaultdict
import sys

input = sys.stdin.readline

T = int(input())

n = int(input())
A = list(map(int, input().split()))
# 구간합 구하기
A_sum = [0] * n
A_sum[0] = A[0]
for i in range(1, n):
    A_sum[i] = A[i] + A_sum[i-1]
    
m = int(input())
B = list(map(int, input().split()))
# 구간합 구하기
B_sum = [0] * m
B_sum[0] = B[0]
for i in range(1, m):
    B_sum[i] = B[i] + B_sum[i-1]
    
# A와 B의 전체 부배열의 합을 구한다.
A_sub = defaultdict(int)
for i in range(n):
    for j in range(i, n):
        A_sub[A_sum[j] - A_sum[i-1] if i != 0 else A_sum[j]] += 1 
        
B_sub = defaultdict(int)
for i in range(m):
    for j in range(i, m):
        B_sub[B_sum[j] - B_sum[i-1] if i != 0 else B_sum[j]] += 1
        
A_sub_sort = sorted(A_sub) # A의 부배열 합 오름차순 정렬
B_sub_sort = sorted(B_sub, reverse=True) # B의 부배열 합 내림차순 정렬

a, b = 0, 0 # A와 B의 부배열 합을 투포인터로 탐색하기 위함
cnt = 0

while a < len(A_sub) and b < len(B_sub):
    section_sum = A_sub_sort[a] + B_sub_sort[b]
    
    if section_sum > T:
        b += 1
    elif section_sum < T:
        a += 1
    else: # 정답이라면
        cnt += A_sub[A_sub_sort[a]] * B_sub[B_sub_sort[b]] # 해당하는 수의 각각의 갯수들끼리 곱해서 더하기 ex) A에서 합이 2가 되는게 3개, B에서 합이 3이 되는게 2개면, 전체 합이 5가 되는건 3 * 2 = 6 개니까
        a += 1 # 뭘 증가시키던 상관없다. 어차피 숫자의 종류로 우리는 접근하기 때문에, 어떤걸 증가시키던 다음 순회에 위의 조건들을 다시 돌게 될거다.
        # b += 1 
            
print(cnt)