파이썬으로 리스트의 모든 연속 부분 배열 최솟값 합 구하기 – 단조 스택 활용법

문제 이해하기

숫자로 이루어진 리스트 nums가 주어졌다고 가정해 보겠습니다. 목표는 연속된 모든 부분 배열(서브리스트)에 대해 각각의 최솟값을 구한 뒤, 이 값들을 모두 더한 합계를 계산하는 것입니다. 답이 매우 커질 수 있으므로, 최종 결과는 10^9 + 7로 나눈 나머지를 반환합니다.

예를 들어 입력이 nums = [5, 10, 20, 10, 0]이라면 출력은 90입니다. 생성 가능한 부분 배열은 총 15개([[5], [10], [20], [10], [0], [5,10], [10,20], [20,10], [10,0], [5,10,20], [10,20,10], [20,10,0], [5,10,20,10], [10,20,10,0], [5,10,20,10,0]])이며, 각 부분 배열의 최솟값은 [5, 10, 20, 10, 0, 5, 10, 10, 0, 5, 10, 0, 5, 0, 0]입니다. 이 값들을 모두 더하면 90이 됩니다.

접근 방법: 단조 스택(Monotonic Stack)

배열의 길이가 n일 때 부분 배열의 개수는 n×(n+1)/2개입니다. 모든 부분 배열을 하나씩 검사하는 브루트포스 방식은 비효율적이지만, 단조 스택 기법을 활용하면 O(n) 시간 안에 문제를 해결할 수 있습니다.

핵심 아이디어는 다음과 같습니다.

• temp_sum에는 현재 위치에서 끝나는 모든 부분 배열의 최솟값 합계를 유지합니다.
• 새로 들어온 값이 스택 꼭대기의 값보다 작거나 같으면, 기존 원소들은 더 이상 최솟값 역할을 할 수 없으므로 스택에서 제거하고 temp_sum에서 그 기여분을 차감합니다.
• 스택의 각 원소는 [인덱스, 값, 기여도] 형태로 저장되며, 기여도는 "이 원소가 최솟값이 되는 부분 배열의 개수 × 값"입니다.
• 매 반복마다 ans에 temp_sum을 더해 전체 합계를 누적합니다.

구체적인 알고리즘 단계는 아래와 같습니다.

1. ans := 0, s := [], temp_sum := 0으로 초기화합니다.
2. nums의 각 인덱스와 값을 순회하며 다음을 수행합니다.
   • 스택이 비어 있지 않고 value ≤ 스택 마지막 원소의 값이면, temp_sum에서 마지막 원소의 기여도를 빼고 스택에서 제거(pop)합니다.
   • 스택이 비어 있다면 [index, value, (index + 1) × value]를 삽입합니다.
   • 그렇지 않다면 [index, value, (index − 스택 마지막 원소의 인덱스) × value]를 삽입합니다.
   • temp_sum에 방금 삽입한 원소의 기여도를 더합니다.
   • ans에 temp_sum을 더합니다.
3. ans mod (10^9 + 7)을 반환합니다.

구현 예제

아래 파이썬 코드를 통해 실제 동작 과정을 확인할 수 있습니다.

def solve(nums):
    ans = 0
    s = []
    temp_sum = 0
    for index, value in enumerate(nums):
        while s and value <= s[-1][1]:
            temp_sum -= s[-1][2]
            s.pop()
        if not s:
            s.append([index, value, (index + 1) * value])
        else:
            s.append([index, value, (index - s[-1][0]) * value])
        temp_sum += s[-1][2]
        ans += temp_sum
    return ans % (10**9 + 7)

nums = [5, 10, 20, 10, 0]
print(solve(nums))

실행 결과

입력:

[5, 10, 20, 10, 0]

출력:

90

복잡도 분석

• 시간 복잡도: O(n) — 각 원소가 스택에 최대 한 번 삽입되고 한 번 제거됩니다.
• 공간 복잡도: O(n) — 최악의 경우 스택에 모든 원소가 저장될 수 있습니다.