スポンサーリンク

907. Sum of Subarray Minimums

mediumとは思えない難しさだったので備忘録。

Sum of Subarray Minimums - LeetCode
Sum of Subarray Minimums - Given an array of integers arr, find the sum of min(b), where b ranges over every (contiguous) subarray of arr. Since the answer may ...

与えられた配列からすべての部分列を作り、各部分列の最小値の合計を求める問題。

Input: arr = [3,1,2,4]
Output: 17
Explanation:
Subarrays are [3], [1], [2], [4], [3,1], [1,2], [2,4], [3,1,2], [1,2,4], [3,1,2,4].
Minimums are 3, 1, 2, 4, 1, 1, 2, 1, 1, 1.
Sum is 17.

初手として全探索が思いつくが、N<=3*10^4のため、O(N^2)となり間に合わない。

以下はTLEになる全探索例。

class Solution:
    def sumSubarrayMins(self, arr: List[int]) -> int:
        mod = 10 ** 9 + 7

        n = len(arr)
        ans = 0
        for start in range(n):
            for length in range(n):
                end = start + length + 1
                if end == n + 1:
                    break

                sub = arr[start:end]
                ans += min(sub)
        
        return ans % mod

計算量を落とすためには発想を変える必要がある。

「部分列の最小値」を求めるのではなく、「ある値が最小となるような部分列」を求めることにする。

サンプルケースの場合で考えると、

3が最小値になるような部分列は[3]

1が最小になるような部分列は[1],[3,1], [1,2],[3,1,2], [1,2,4], [3,1,2,4]

2が最小値になるような部分列は[2],[2,4]

4が最小値になるような部分列は[4]

となる。

ある値xが最小値となるような部分列を求めるには、xの左側と右側にxよりも小さい値がなければいい。

これはmonotonic stackというデータ構造を使えば求めることが出来る。

Introduction to Monotonic Stack - GeeksforGeeks
A Computer Science portal for geeks. It contains well written, well thought and well explained computer science and programming articles, quizzes and practice/c...
Monotonic Stack - Algorithm

以上の考察から、問題を解くアプローチは

1.値xより大きい値が左側に何個あるか数えてLとする。これはmonotonic stackで数えられる。

2.同様に、右側にある大きい値の個数をRとする。

3.1と2でxが最小となるような部分列の両端が求まるので、xが最小となる部分列の個数は、左端を0〜L、右端を0〜Rの選び方で決まるのでLxR個となる。

となる。

class Solution:
    def sumSubarrayMins(self, arr: List[int]) -> int:
        mod = 10 ** 9 + 7

        n = len(arr)
        
        left = [0] * n
        stack_left = []
        for i in range(n):
            count = 1
            while stack_left and stack_left[-1][0] >= arr[i]:
                count += stack_left[-1][1]
                stack_left.pop()
            
            left[i] = count
            stack_left.append((arr[i], count))
        
        right = [0] * n
        stack_right = []
        for i in range(n-1, -1, -1):
            count = 1
            while stack_right and stack_right[-1][0] > arr[i]:
                count += stack_right[-1][1]
                stack_right.pop()
            
            right[i] = count
            stack_right.append((arr[i], count))
        
        ans = 0
        for i in range(n):
            ans += arr[i] * left[i] * right[i]
        
        return ans % mod

monotonic stackを左右で2回、答えを求める計算が1回なので、計算量はO(3xN)=O(N)

メモリは左右の個数とスタックの記録が必要になるのでO(4xN)=O(N)となる。



コメント

タイトルとURLをコピーしました