[Python] 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] 

https://www.acmicpc.net/problem/2143

 

 

 

 

난이도

골드 3

 

 

 

 

 

 

 

내 코드

import sys

t = int(sys.stdin.readline())
n = int(sys.stdin.readline())
A = list(map(int,sys.stdin.readline().split()))
m = int(sys.stdin.readline())
B = list(map(int,sys.stdin.readline().split()))
answer = 0

sumA = [0]
sumB = [0]
for i in range(n) :
    sumA.append(sumA[i] + A[i])
for i in range(m) :
    sumB.append(sumB[i] + B[i])

cnt_A = dict()
cnt_B = dict()
for i in range(len(sumA)-1) :
    for j in range(i+1,len(sumA)) :
        tmp = sumA[j] - sumA[i]
        if tmp in cnt_A :
            cnt_A[tmp] += 1
        else :
            cnt_A[tmp] = 1
for i in range(len(sumB)-1) :
    for j in range(i+1,len(sumB)) :
        tmp = sumB[j] - sumB[i]
        if tmp in cnt_B :
            cnt_B[tmp] += 1
        else :
            cnt_B[tmp] = 1

for a in cnt_A :
    if t-a in cnt_B :
        answer += (cnt_A[a] * cnt_B[t-a])

print(answer)

 

 

 

 

풀이과정

두 배열의 누적합 배열을 만들고, 그 누적합 배열을 이용해서 A와 B 배열에서 나올 수 있는 수들을 각각의 딕셔너리에 저장했다.

딕셔너리의 key 값은 나올 수 있는 수로 하고 value 값은 나온 횟수를 카운트 하도록 했다.

 

이 후 처음에는 A,B 배열을 오름차순으로 정렬하여 이분탐색으로 풀었더니 시간초과가 떴다.

 

생각해보니 그냥 딕셔너리의 특징을 이용해서 세 줄만에 가볍게 풀 수 있었다.