여러 개의 스택이 리스트 형태로 주어져 있다고 가정해 보겠습니다. 우리는 임의의 스택에서 원하는 만큼 요소를 팝(pop)해 제거할 수 있으며, 이때 모든 스택의 합계가 서로 같아지도록 만들면서 얻을 수 있는 최대 합계를 구해야 합니다.
문제 예시
예를 들어 입력이 다음과 같다고 해봅시다.
stacks = [[3, 4, 5, 6], [5, 6, 1, 4, 4], [10, 2, 2, 2]]
이 경우 정답은 12입니다. 아래와 같은 연산을 수행하면 세 스택 모두 합계를 12로 맞출 수 있습니다.
첫 번째 스택에서 [6]을 팝 → 남은 스택 [3, 4, 5], 합계 = 12
두 번째 스택에서 [4, 4]를 팝 → 남은 스택 [5, 6, 1], 합계 = 12
세 번째 스택에서 [2, 2]를 팝 → 남은 스택 [10, 2], 합계 = 12
접근 방법: 접두사 합(Prefix Sum) 활용
핵심 아이디어는 간단합니다. 스택은 위(top)에서부터 요소를 제거할 수 있으므로, 각 스택에서 만들어질 수 있는 합계는 앞에서부터 차례대로 누적한 접두사 합들입니다. 따라서 모든 스택에 공통으로 나타나는 접두사 합 중 가장 큰 값을 찾으면 그것이 곧 정답이 됩니다.
알고리즘 단계
sums := 빈 딕셔너리(맵)를 생성합니다.
각 스택 stk에 대해 다음을 반복합니다.
s := 0으로 초기화합니다.
스택의 각 요소 n에 대해 다음을 수행합니다.
s := s + n (누적 합 계산)
sums[s] := sums[s] + 1 (해당 합계가 등장한 스택의 수 기록)
ans := 0으로 초기화합니다. (모든 스택을 전부 비우면 합계 0은 항상 가능하므로)
sums의 각 키-값 쌍 (s, f)에 대해 다음을 검사합니다.
f가 스택 개수 이상이면(즉, 모든 스택에서 등장하는 합계이면) 그리고 s > ans이면 ans := s로 갱신합니다.
ans를 반환합니다.
파이썬 구현 예제
아래 구현을 통해 더 잘 이해해 보겠습니다.
from collections import defaultdict
class Solution:
def solve(self, stacks):
sums = defaultdict(int)
for stk in stacks:
s = 0
for n in stk:
s += n
sums[s] += 1
ans = 0
for s, f in sums.items():
if f >= len(stacks) and s > ans:
ans = s
return ans
ob1 = Solution()
stacks = [
[3, 4, 5, 6],
[5, 6, 1, 4, 4],
[10, 2, 2, 2]
]
print(ob1.solve(stacks))
입력
stacks = [[3, 4, 5, 6], [5, 6, 1, 4, 4], [10, 2, 2, 2]]
출력
12
복잡도 분석
시간 복잡도는 O(N)입니다. 여기서 N은 모든 스택에 담긴 요소의 총 개수로, 각 요소를 한 번씩만 방문하면 충분하기 때문입니다. 공간 복잡도 역시 저장해야 하는 고유한 접두사 합의 개수에 비례하므로 최대 O(N)입니다.