leetcode : Sum of Subarray Ranges

Sum of Subarray Ranges - LeetCode

Sum of Subarray Ranges - LeetCode

 

문제를 해석해보면 주어진 array에서 임의의 i,j에 대해( 0 <= i <= j <= array.len-1) max(array[i:j]) - min(array[i:j])의 합을 구하라는 것이다. (i = j인 경우는 무조건 0이 나오니까 i < j 조건이 사실상 맞긴 하다.)

1, 2, ,3이 주어지면

(1,3) -> 2

(1,2) -> 1

(2,3) -> 1

합은 4

이런식이다.

100% naive하게 생각하면  O(N^3)의 시간복잡도가 나올 것이다. 각각의 원소에 대해 그 원소보다 오른쪽에 있는 원소들중 큰 것을 찾아 차를 구해줘야되기 때문이다.

array를 첫 번째 원소부터 마지막까지 스캔하면서 각 시점에 특정 조건의 '합'을 실시간으로 구할 수 있을까? range를 구하기 위해서는 최대값과 최소값을 알고있어야 한다.

monotonic stack(monotone stack | JusticeHui가 PS하는 블로그)을 이용한다면.. 인덱스 i를 볼때 array[0]부터 array[i]까지의 요소들 중 array[i]보다 큰값이나(내림차순) 작은값들을(오름차순) 빠르게 알 수 있다. 최대값, 최소값도 당연히 알 수 있다. 근데 어케해야될지 모르겠음 몇시간째 고민중인데.

 

Java | Well explained with examples | O(n) - LeetCode Discuss

Java | Well explained with examples | O(n) - LeetCode Discuss

와우.. 일단 기본아이디어는 이거.. 특정 원소가 min값인 걸 기준으로 잡는거랑 양쪽으로 구해서 곱하는거 생각은 했는데 (max-min) + (max-min) + ... 를 max끼리 묶고 min끼리 묶는건 생각도 못했네.

 

min stack과 max stack을 만들어 2스캔으로 풀었다.

pop하는 시점에 우리가값이 계산된다.

 

class Solution:
    def subArrayRanges(self, nums: List[int]) -> int:
        answer = 0
        
        stack = []
        i = 0
        for i in range(len(nums) + 1):
            v = nums[i] if i < len(nums) else -10**9-1
            while stack and nums[stack[-1]] >= v:
                ti = stack.pop()
                tv = nums[ti]
                
                k = stack[-1] if stack else -1 
                
                answer -= (i - ti) * (ti - k) * tv
            stack.append(i)
            
        stack = []
        for i in range(len(nums) + 1):
            v = nums[i] if i < len(nums) else 10**9+1
            while stack and nums[stack[-1]] <= v:
                ti = stack.pop()
                tv = nums[ti]
                
                k = stack[-1] if stack else -1
                
                answer += (i - ti) * (ti - k) * tv
            stack.append(i)
        
        return answer

 

공통부분을 함수로 묶어서 함수(max) - 함수(min)으로 구해도 된다.

long long subArrayRanges(vector<int>& n) {
    return sumSubarrayComp(n, less<int>()) - sumSubarrayComp(n, greater<int>());
}    
long long sumSubarrayComp(vector<int>& n, function<bool (int, int)> comp) {
    long long res = 0;
    vector<int> s;
    for (int i = 0; i <= n.size(); ++i) {
        while (!s.empty() && (i == n.size() || comp(n[s.back()], n[i]))) {
            int j = s.back(), k = s.size() < 2 ? -1 : s[s.size() - 2]; 
            res += (long long)(i - j) * (j - k) * n[j];
            s.pop_back();
        }
        s.push_back(i);
    }    
    return res;
}

 

이문제.. 재밌었다. 몇시간을 고민했네 진짜