백준 2143번: 두 배열의 합 (Python)

문제

한 배열 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]

 

해설

이분탐색과 누적합 알고리즘을 이용해야 하는 문제다. 태그를 보지 않고 풀었더니 삽질을 너무 많이했다. 그래도 의미있는 실패였다고 생각하기에 나중에 기억하기 위해 풀이를 올려본다.

 

💡 사고의 흐름 💡 

1. 첫 번째 시도

"백트래킹+브루트포스를 활용해 리스트 A, 리스트 B에서 나올 수 있는 모든 합의 경우의 수를 딕셔너리에 저장한 다음, 두 딕셔너리를 비교하여 합이 S가 되는 경우를 찾으면 되지 않을까?"

 

=> 결과: 근본적으로 문제를 잘못 이해함. '부 배열의 합'은 연속된 부분수열의 합을 의미함. 백트래킹으로 모든 조합을 구하는 풀이는 틀린 접근이었음.

(ex. {1, 3, 1, 2} 에서 {1, 3, 1, 2} 처럼 건너뛰어서 합을 구할 수 없음)

게다가 입력값을 고려할 때 브루트포스를 사용했으면 시간초과를 받았을 것으로 예상.

 

2. 두 번째 시도

"그렇다면 연속된 부분 배열의 합을 구하는 것은 투포인터를 활용하면 쉽게 구할 수 있으니 투 포인터를 이용해야겠다."

 

=> 결과: 입력값에 음수가 들어오는 경우 투포인터를 이용한 부분합 구하기 알고리즘을 활용할 수 없음. 다시 코드 갈아엎음.

 

3. 마지막 시도

"우선 구간합을 빠르게 구하기 위해 리스트 A, 리스트 B의 누적합 테이블을 구하고, 가능한 (i, j) 쌍을 모두 대입해가며 i~j 구간의 구간합을 각 딕셔너리에 key 값으로 저장해보자. 이후 A, B 두 딕셔너리의 key 값을 비교해가며 S가 되는 경우를 찾아보자."

 

 

❗ 풀이 과정 ❗

 

1. 누적합 테이블 구하기

 

[2, 1, 5, 7, 3] 이라는 배열이 있다고 해보자. 만약 인덱스 3~5 구간의 합을 구하고 싶다면 어떻게 해야 할까?

시작 인덱스부터 끝 인덱스까지 더해나가는 방식은 최악의 경우 O(n^2) 의 시간복잡도를 갖기 때문에, 누적합 테이블을 만들어 O(n)에 해결해주는 것이 좋다. 

 

누적합 테이블은 다음과 같이 생겼다. 0~현재 인덱스까지의 누적 합을 저장하는 테이블이다.

누적합 테이블

 

def GeneratePrefixSum(arr):
    size = len(arr)
    DP = [0] * size
    DP[0] = arr[0]
    for i in range(1, size):
        DP[i] = DP[i-1] + arr[i]
    return DP

lst = [2,1,5,7,3]
print(GeneratePrefixSum(lst))
# result: [2, 3, 8, 15, 18]

 

2. 누적합 테이블을 이용해 구간합 구하기

구간합의 시작 인덱스를 s, 끝 인덱스를 e라고 한다면 s~e 구간의 합은 위 그림과 같이 구할 수 있다.

구간합[s:e] = 누적합[e] - 누적합[s-1] 

따라서 만들어진 누적합 테이블을 활용해 구간합을 구하는 함수는 다음과 같이 만들 수 있다.

# 누적합 테이블 만드는 함수
def GeneratePrefixSum(arr):
    size = len(arr)
    DP = [0] * size
    DP[0] = arr[0]
    for i in range(1, size):
        DP[i] = DP[i-1] + arr[i]
    return DP

# start ~ end 사이의 구간합 구하는 함수
def GetPrefixSum(DP, start, end):
    DP = [0] + DP  # margin
    return DP[end] - DP[start-1]


table = GeneratePrefixSum([2, 1, 5, 7, 3])
print(GetPrefixSum(table, 3, 5))
#result: 15

 

3. 모든 구간합 경우의 수를 구해 딕셔너리에 카운트하기

 

앞선 과정을 통해 빠르게 구간합을 구할 수 있게 되었다. 이제 i와 j를 증가시키면서 리스트 A, 리스트 B의 모든 구간합을 체크해야 한다. 

 

리스트 [2, 1, 5, 7, 3] 을 예로 든다고 하면.

 

A [ 0 ~ 0 ] = 2

A [ 0 ~ 1 ] = 3

A [ 0 ~ 2 ] = 8

...

A [ 3 ~ 5 ] = 15

A [ 4 ~ 5 ] = 10

A [ 5  ~ 5 ] = 3

.이런식으로 모든 구간합을 구할 수 있을 것이다. 이렇게 나온 구간합을 딕셔너리의 key 값으로 추가하면 다음과 같은 결과를 얻을 수 있을 것이다.

{2: 1, 3: 2, 8: 1, 15: 2, 18: 1, 1: 1, 6: 1, 13: 1, 16: 1, 5: 1, 12: 1, 7: 1, 10: 1}

그러면 딕셔너리의 key 값은 구간합의 값, value 값은 구간합이 나온 횟수를 의미하게 된다.

 

리스트 A, 리스트 B 모두 위와 같은 딕셔너리를 구해준다.

 

4. 딕셔너리의 key 값을 비교해가며 S가 되는 경우 찾기

 

거의 다 왔다. 우리는 각 리스트 A, B 에서 나온 구간합의 set과 나온 횟수를 저장하는 딕셔너리를 갖고 있다. 이제 여기서 우리는 다음과 같이 생각할 수 있다.

 

dict_a 의 key를 모두 훑는다고 했을 때, (S - a.key) 가 dict_b.keys() 안에 존재한다면 S를 만들 수 있는 경우의 수가 존재한다는 것이다. 이 경우의 수는 ( dict_a.value * dict_b.value ) 로 구할 수 있을 것이다. (조합)

 

 

 

A = [1,2,3], B = [1,1,2]를 예시로 든다면 다음과 같다.

 

이렇게 모든 key를 서치하면서 S가 되는 경우의 수를 count에 누적해 더해나간다면 결과를 얻을 수 있다.

 

코드

# 누적합 테이블 만드는 함수
def GeneratePrefixSum(arr):
    size = len(arr)
    DP = [0] * size
    DP[0] = arr[0]
    for i in range(1, size):
        DP[i] = DP[i-1] + arr[i]
    return DP

# start ~ end 사이의 구간합 구하는 함수
def GetPrefixSum(DP, start, end):
    DP = [0] + DP  # margin
    return DP[end] - DP[start-1]


T = int(input())
n = int(input())
A = list(map(int, input().split()))
m = int(input())
B = list(map(int, input().split()))

DP_A = GeneratePrefixSum(A)
DP_B = GeneratePrefixSum(B)
DICT_A = dict()
DICT_B = dict()

for i in range(1, n+1):
    for j in range(i, n+1):
        res = GetPrefixSum(DP_A, i, j)
        if res in DICT_A.keys():
            DICT_A[res] += 1
        else:
            DICT_A[res] = 1

for i in range(1, m+1):
    for j in range(i, m+1):
        res = GetPrefixSum(DP_B, i, j)
        if res in DICT_B.keys():
            DICT_B[res] += 1
        else:
            DICT_B[res] = 1

count = 0
for sum_a in DICT_A.keys():
    if (T-sum_a) in DICT_B.keys():
        count += DICT_A[sum_a] * DICT_B[T-sum_a]
print(count)

여담

골드 3치고는 꽤 어려운 편이라고 생각이 들었다. 특히 마지막 접근은 골드1 문제인 [부분수열의 합 2] 문제 풀이에서 영감을 많이 받았다. 역시 알고리즘 문제는 많이 풀어야 실력이 느는 것 같다. 

백준 1208번: 부분수열의 합 2