[Python/파이썬] 백준 2143번 두 배열의 합

2143번: 두 배열의 합

문제

한 배열 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(=5) = A[1] + B[1] + B[2]
      = A[1] + A[2] + B[1]
      = A[2] + B[3]
      = A[2] + A[3] + B[1]
      = A[3] + B[1] + B[2]
      = A[3] + A[4] + B[3]
      = A[4] + B[2] 

입력

첫째 줄에 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을 출력한다.

코드

• 투포인터 시도(안됨)

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

a_move = False

a_start, a_end, b_start, b_end = 0, 0, 0, 0
res = 0
while a_end < n - 1 and b_end < m - 1:
    now_sum = sum(a[a_start:a_end+1]) + sum(b[b_start:b_end + 1])
    if now_sum == t:
        res += 1
       # a가 움직일 차례
        if not a_move:
            a_move = True
            if a_end < n - 1:
                a_start += 1
                a_end += 1
            else:
                break
        # b가 움직일 차례
        else:
            a_move = False
            if b_end < m - 1:
                b_start += 1
                b_end += 1
            else:
                break
    elif now_sum < t:
        # a가 움직일 차례
        if not a_move:
            a_move = True
            if a_end < n - 1:
                a_end += 1        
    # b가 움직일 차례
        else:
            a_move = False
            if b_end < m - 1:
                b_end += 1
    else:
        # a가 움직일 차례
        if not a_move:
            a_move = True
            a_start += 1        
    # b가 움직일 차례
        else:
            a_move = False
            b_start += 1

print(res)

 

• 부분합 + 딕셔너리 이용

from collections import defaultdict
import sys
input = sys.stdin.readline

t = int(input())
n = int(input())
a_arr = list(map(int, input().split()))
m = int(input())
b_arr = list(map(int, input().split()))

# A 배열의 부분합을 딕셔너리에 개수만 저장
a_dict = defaultdict(int)
for i in range(n):
    for j in range(i, n):
        a_dict[sum(a_arr[i:j+1])] += 1
# B 배열의 부분합은 구하는 동시에 t에서 그 값을 빼서 A 부분합 딕셔너리 인덱스로 사용
res = 0
for i in range(m):
    for j in range(i, m):
        res += a_dict[t - sum(b_arr[i:j+1])]
    
print(res)

처음에는 A배열의 부분합과 B배열의 부분합 개수를 저장한 딕셔너리를 따로 만들었는데 그랬더니 메모리 초과가 발생해서 B 배열의 부분합은 구하는 동시에 t에서 그 값을 빼서 a_dict의 인덱스로 사용하였다. 그 인덱스를 이용하여 a_dict의 밸류값들을 더해주면 답을 구할 수 있다.

 

아래의 블로그 참고했음...이중탐색은 사용하지 않고 파이썬의 자료구조를 잘 활용한 풀이였다.

[알고리즘][python][백준 2143]두 배열의 합