mediumとは思えない難しさだったので備忘録。
https://leetcode.com/problems/sum-of-subarray-minimums/description/
与えられた配列からすべての部分列を作り、各部分列の最小値の合計を求める問題。
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]
１が最小になるような部分列は[1],[3,1], [1,2],[3,1,2], [1,2,4], [3,1,2,4]
２が最小値になるような部分列は[2],[2,4]
４が最小値になるような部分列は[4]
となる。
ある値ｘが最小値となるような部分列を求めるには、ｘの左側と右側にｘよりも小さい値がなければいい。
これはmonotonic stackというデータ構造を使えば求めることが出来る。
https://www.geeksforgeeks.org/introduction-to-monotonic-stack/
https://liuzhenglaichn.gitbook.io/algorithm/monotonic-stack
以上の考察から、問題を解くアプローチは
１．値ｘより大きい値が左側に何個あるか数えてLとする。これはmonotonic stackで数えられる。
２．同様に、右側にある大きい値の個数をRとする。
３．１と２でｘが最小となるような部分列の両端が求まるので、ｘが最小となる部分列の個数は、左端を０〜L、右端を０〜Rの選び方で決まるのでLｘR個となる。
となる。
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 % modmonotonic stackを左右で２回、答えを求める計算が１回なので、計算量はO(3xN)=O(N)
メモリは左右の個数とスタックの記録が必要になるのでO(4xN)=O(N)となる。
コメント