문제 개요
여러 개의 스택으로 구성된 목록과 정수 k가 주어졌을 때, 임의의 스택 조합에서 정확히 k개의 요소를 팝(pop)하여 만들 수 있는 최대 합을 구하는 프로그램을 작성해야 합니다. 스택은 LIFO(후입선출) 구조이므로, 각 스택에서는 반드시 맨 위 요소부터 순서대로 꺼내야 한다는 점이 핵심입니다.
예를 들어 입력이 다음과 같다고 가정해 보겠습니다.
stacks = [[50, -4, -15], [2], [6, 7, 8]], k = 4
이 경우 출력은 39입니다. 첫 번째 스택의 세 요소(-15, -4, 50)를 모두 팝하고, 마지막 스택의 맨 위 요소(8)를 하나 더 팝하면 -15 + (-4) + 50 + 8 = 39라는 합을 얻을 수 있기 때문입니다.
풀이 접근 방법
이 문제는 재귀 호출을 통해 가능한 모든 조합을 탐색하는 방식으로 해결할 수 있습니다. rec(i, n) 함수를 정의하고 아래 단계에 따라 진행합니다.
rec(i, n): i번째 스택부터 탐색을 시작하며, 지금까지 팝한 요소의 총 개수는 n입니다.
n이 k와 같으면 더 이상 팝할 요소가 없으므로 0을 반환합니다.
n이 k보다 크면 조건을 초과했으므로 음의 무한대(-inf)를 반환합니다.
i가 스택의 총 개수와 같으면 더 이상 사용할 스택이 없으므로 음의 무한대를 반환합니다.
i가 마지막 스택의 인덱스라면 남은 개수 needed = k - n을 계산합니다. needed가 현재 스택의 크기보다 크면 음의 무한대를 반환하고, 그렇지 않으면 스택의 뒤쪽 needed개 요소의 합을 반환합니다.
그 외의 경우에는 현재 스택에서 팝할 개수를 하나씩 늘려 가며 누적합 su를 갱신하고, rec(i + 1, ...)의 반환값과 더한 localres 중 최댓값을 res에 저장합니다.
마지막으로 현재 스택에서 아무것도 팝하지 않는 경우 rec(i + 1, n)과 비교해 더 큰 값을 반환합니다.
메인 메서드에서는 rec(0, 0)을 호출해 최종 결과를 구합니다.
Python 구현 예제
import math class Solution: def solve(self, stacks, k): def rec(i, n): if n == k: return 0 if n > k: return -math.inf if i == len(stacks): return -math.inf if i == len(stacks) - 1: needed = k - n if needed > len(stacks[i]): return -math.inf else: return sum(stacks[i][-needed:]) res, su = -math.inf, 0 for sti in range(len(stacks[i]) - 1, -1, -1): su += stacks[i][sti] localres = su + rec(i + 1, n + len(stacks[i]) - sti) res = max(res, localres) return max(res, rec(i + 1, n)) return rec(0, 0) ob = Solution() stacks = [ [50, -4, -15], [2], [6, 7, 8] ] k = 4 print(ob.solve(stacks, k))
입력
[[50, -4, -15], [2], [6, 7, 8]], 4
출력
39
동작 원리 살펴보기
rec 함수는 각 스택에 대해 두 가지 선택지를 고려합니다. 첫째, 현재 스택에서 하나 이상의 요소를 팝하는 경우입니다. 이때 맨 위 요소부터 순서대로 꺼내야 하므로 스택의 끝에서부터 거꾸로 순회하며 누적합을 계산합니다. 둘째, 현재 스택을 건드리지 않고 다음 스택으로 넘어가는 경우입니다. 두 결과 중 더 큰 값을 선택하며, 마지막 스택에 도달하면 남은 개수만큼 정확히 팝할 수 있는지만 확인하면 됩니다.
이 방식은 모든 유효한 조합을 빠짐없이 탐색하면서도, 불필요한 분기(n > k 또는 스택의 요소 부족)는 즉시 -inf로 잘라내기 때문에 효율적으로 동작합니다. 참고로 스택의 개수와 크기가 커지는 경우에는 메모이제이션(memoization)을 적용해 동일한 (i, n) 상태의 결과를 캐싱함으로써 실행 시간을 크게 단축할 수 있습니다.