파이썬으로 두 배열에서 합이 같은 부분 배열 찾기

문제 정의

크기가 N인 두 배열 P와 Q가 있고, 두 배열에는 1부터 N 사이의 숫자들이 들어 있다고 가정해 봅시다. 이때 각 배열에서 연속된 구간, 즉 부분 배열(sub-array)을 하나씩 골라 두 합이 정확히 같아지도록 만들어야 합니다. 조건을 만족하는 부분 배열의 인덱스를 반환하고, 가능한 조합이 존재하지 않는다면 -1을 출력하면 됩니다.

예를 들어 입력이 다음과 같다고 해보겠습니다.

P = [2, 3, 4, 5, 6]
Q = [9, 3, 2, 6, 5]

이 경우 첫 번째 배열의 인덱스 0, 1, 2와 두 번째 배열의 인덱스 0이 정답입니다. 실제로 P[0..2]의 합은 2 + 3 + 4 = 9이고, Q[0] = 9이므로 두 부분 배열의 합이 서로 같습니다.

풀이 전략

핵심 아이디어는 누적 합(prefix sum)과 차이 값 추적입니다. 먼저 두 배열을 누적 합 형태로 변환하면 특정 구간의 합은 두 누적 값의 차이로 표현할 수 있습니다. 따라서 두 배열의 누적 합 차이가 동일한 지점이 다시 등장하면, 그 사이 구간의 합이 반드시 같다는 사실을 이용해 답을 찾을 수 있습니다.

get_subarray() 함수의 동작 과정

함수는 P, Q, swap 세 개의 매개변수를 받습니다.
N := P의 크기
index := 새로운 딕셔너리(맵)
difference := 0, j := 0으로 초기화
index[0] := (-1, -1)을 저장합니다. 처음 위치부터 시작하는 구간도 처리하기 위함입니다.
i를 0부터 N-1까지 반복합니다.
- Q[j] < P[i]인 동안 j를 1씩 증가시킵니다.
- difference := Q[j] - P[i]
- difference가 index에 이미 존재하면:
  - swap이 참인 경우: idx := index[Q[j] - P[i]]를 꺼낸 뒤, 첫 번째 배열에는 idx[1]+1부터 j까지, 두 번째 배열에는 idx[0]+1부터 i까지의 인덱스를 출력합니다.
  - 그렇지 않은 경우: idx := index[Q[j] - P[i]]를 꺼낸 뒤, 첫 번째 배열에는 idx[0]+1부터 i까지, 두 번째 배열에는 idx[1]+1부터 j까지의 인덱스를 출력합니다.
  결과를 출력한 후 함수를 종료합니다.
- 같은 차이 값이 없었다면 index[difference] := (i, j)를 저장합니다.
반복이 끝날 때까지 답을 찾지 못했다면 -1을 출력합니다.

메인 로직

cumsum() 함수를 이용해 P와 Q를 누적 합 배열로 변환합니다.
N := P의 크기
만약 Q[N - 1] > P[N - 1]이면 get_subarray(P, Q, False)를 호출합니다.
그렇지 않으면 get_subarray(Q, P, True)를 호출합니다.

구현 예제

아래 파이썬 코드를 통해 더 자세히 이해해 보겠습니다.

def show_res(x, y, num):
    print("Indices of array", num, ":", end=" ")
    for i in range(x, y):
        print(i, end=", ")
    print(y)

def get_subarray(P, Q, swap):
    N = len(P)
    index = {}
    difference, j = 0, 0
    index[0] = (-1, -1)
    for i in range(0, N):
        while Q[j] < P[i]:
            j += 1
        difference = Q[j] - P[i]
        if difference in index:
            if swap:
                idx = index[Q[j] - P[i]]
                show_res(idx[1] + 1, j, 1)
                show_res(idx[0] + 1, i, 2)
            else:
                idx = index[Q[j] - P[i]]
                show_res(idx[0] + 1, i, 1)
                show_res(idx[1] + 1, j, 2)
            return
        index[difference] = (i, j)
    print(-1)

def cumsum(arr):
    n = len(arr)
    for i in range(1, n):
        arr[i] += arr[i - 1]

P = [2, 3, 4, 5, 6]
Q = [9, 3, 2, 6, 5]
cumsum(P)
cumsum(Q)
N = len(P)
if Q[N - 1] > P[N - 1]:
    get_subarray(P, Q, False)
else:
    get_subarray(Q, P, True)

입력

[2, 3, 4, 5, 6],[9, 3, 2, 6, 5]

출력

Indices of array 1 : 0, 1, 2
Indices of array 2 : 0

복잡도 분석

포인터 i와 j가 항상 앞으로만 이동하므로 전체 순회 횟수는 선형적으로 제한됩니다. 따라서 시간 복잡도는 O(N)이며, 차이 값을 저장하는 딕셔너리의 크기도 최대 N이므로 공간 복잡도 역시 O(N)입니다.