一道OA题，单调栈，想了半天O(n)解没什么思路，求大家帮忙看看。。。

我之前做的，c++的代碼，你們正好幫我看看有沒有bug。謝謝啦！

    // The idea here is to find the minimum set that n can contribute, and then expand on both sides to calculate
    // its total contribution, and sum them up.
    // idxM1[i]: for nums[i], the index of nums[i] - 1 <= it bounds one side the min-set can expand
    // lSm[i]: for nums[i], the biggest index with number less than nums[i] - 1
    // rSm[i]: for nums[i], the smallest index with number less than nums[i] - 1
    // illustration:

    /* idx  0  1  2  3  4  5
       num  2  4  3  1  6  5
     idxM1  3  2  0 -1  5  1    <- for 4, the index for 4-1=3 is 2, for 1, it's the smallest so -1
       lSm -1  0 -1 -1  3  3    <- for 4, the biggest index would be 0 to form min-set [2 4]
       rSm -1  3  3 -1 -1 -1    <- for 4, the smallest index would be 3 to form min-set [4 3 1], of course since 3 is in it
                                   no superset of [4 3 1] can be found
    */
    long long countImbalance(vector<int>& nums) {
        int n = nums.size();
        vector<int> lSm(n, -1), rSm(n, -1), num2idx(n, -1), idxM1(n, -1);
        // num2idx[i] stores the index of i - 1
        // num2idx should be: {3, 0, 2, 1, 5, 4}
        for (int i = 0; i < n; ++i) {
            num2idx[nums[i] - 1] = i;
        }
        // idxM1 should be {3, 2, 0, -1, 5, 1}
        // for nums[i], the index for nums[i] - 1 is idxM1[nums[i] - 1]
        for (int i = 0; i < n; ++i) {
            if (nums[i] == 1) continue;
            idxM1[i] = num2idx[nums[i] - 2];
        }
        // now we should get the closest elem that is next to nums[i], on both sides
        stack<int> stk{};
        // lSm should be: {-1, 0, -1, -1, 3, 3}
        for (int i = 0; i < n; ++i) {
            while (stk.size() && (nums[stk.top()] > nums[i] || nums[stk.top()] == nums[i] - 1)) stk.pop();
            if (stk.size()) lSm[i] = stk.top();
            stk.push(i);
        }
        // rSm should be: {-1, 3, 3, -1, -1}
        while (stk.size()) stk.pop();
        for (int i = n - 1; i >= 0; --i) {
            while (stk.size() && (nums[stk.top()] > nums[i] || nums[stk.top()] == nums[i] - 1)) stk.pop();
            if (stk.size()) rSm[i] = stk.top();
            stk.push(i);
        }
        long long ttl = 0;
        for (int i = 0; i < n; ++i) {
            int rChoice = 0, lChoice = 0;
            if (idxM1[i] < i) {
                if (rSm[i] != -1) {
                    // ... 3 4 1 ... [n=7]
                    rChoice = n - rSm[i];
                    lChoice = i - idxM1[i];
                    ttl += lChoice * rChoice;
                }
                if (lSm[i] > idxM1[i]) {
                    // ... 3 1 4 ... [n=7]
                    lChoice = lSm[i] - idxM1[i];
                    rChoice = n - i;
                    ttl += lChoice * rChoice;
                }
            }
            else {  // idxM1[i] > i
                if (lSm[i] != -1) {
                    // ... 1 4 3 ... [n=7]
                    rChoice = idxM1[i] - i;
                    lChoice = lSm[i] + 1;
                    ttl += lChoice * rChoice;
                }
                if (rSm[i] != -1 && rSm[i] < idxM1[i]) {
                    // ... 1 4 3 ... [n=7]
                    lChoice = i + 1;
                    rChoice = idxM1[i] - rSm[i];
                    ttl += lChoice * rChoice;
                }
            }
        }
        return ttl;
    }

coolwenwen 发表于 2022-9-29 13:18
对于一个数n，要n对imbalance有贡献的话，那么n-1不能和n在同一个数组里。所以你要对每个n找到n-1的位置， ...

是的，根本不是单调栈

https://leetcode.com/playground/55hXJqEY

bydstand 发表于 2022-9-29 18:51
是的，根本不是单调栈

https://leetcode.com/playground/55hXJqEY

思考了半天，终于发现哪里不对了
你可以试一下这个例子：[3,4,6,5,1,2]。答案应该是9，但是贴的代码输出是7。输出一下每个元素的贡献，我贴在下面：
3 4
4 0
6 0
5 3
2 0
很显然，4（或者6）的贡献缺少了[4,6]这个subarray。
这个解法的问题出在，只是单纯考虑n出现在subarray中，但n-1不出现这一情况。但是这个条件是不对的，首先一个问题就是没有考虑n是否为这个subarray中最小的元素。其次就是我上面那个情况，如果有两个个元素和比它小1的元素相邻，那么这两个元素夹在中间的subarray的贡献就直接消失了。

Sibra_Z 发表于 2022-9-29 19:39
思考了半天，终于发现哪里不对了
你可以试一下这个例子：[3,4,6,5,1,2]。答案应该是9，但是贴的代码输出 ...

嗷，有道理，那还得加上判断大小关系

Sibra_Z 发表于 2022-9-29 17:39
思考了半天，终于发现哪里不对了
你可以试一下这个例子：[3,4,6,5,1,2]。答案应该是9，但是贴的代码输出 ...

你說的是對的。我貼錯了我的代碼版本，因為是在code裡面調試，試過好幾種解法。

1.   // its total contribution, and sum them up.
   
2.   // idxM1[i]: for nums[i], the index of nums[i] - 1 <= it bounds one side the min-set can expand
   
3.   // lSm[i]: for nums[i], the biggest index with number less than nums[i] - 1
   
4.   // rSm[i]: for nums[i], the smallest index with number less than nums[i] - 1
   
5.   // illustration:
   
6. 
   
7.   /* idx  0  1  2  3  4  5
   
8.      num  2  4  3  1  6  5
   
9.    idxM1  3  2  0 -1  5  1    <- for 4, the index for 4-1=3 is 2, for 1, it's the smallest so -1
   
10.      lSm -1  0 -1 -1  3  3    <- for 4, the biggest index would be 0 to form min-set [2 4]
    
11.      rSm -1  3  3 -1 -1 -1    <- for 4, the smallest index would be 3 to form min-set [4 3 1], of course since 3 is in it
    
12.                                  
    
13.   */
    
14.   long long countImbalance(vector<int>& nums) {
    
15.       int n = nums.size();
    
16.       vector<int> lSm(n, -1), rSm(n, -1), num2idx(n, -1), idxM1(n, -1);
    
17.       // num2idx[i] stores the index of i - 1
    
18.       // num2idx should be: {3, 0, 2, 1, 5, 4}
    
19.       for (int i = 0; i < n; ++i) {
    
20.           num2idx[nums[i] - 1] = i;
    
21.       }
    
22.       // idxM1 should be {3, 2, 0, -1, 5, 1}
    
23.       // for nums[i], the index for nums[i] - 1 is idxM1[nums[i] - 1]
    
24.       for (int i = 0; i < n; ++i) {
    
25.           if (nums[i] == 1) continue;
    
26.           idxM1[i] = num2idx[nums[i] - 2];
    
27.       }
    
28.       // now we should get the closest elem that is next to nums[i], on both sides
    
29.       stack<int> stk{};
    
30.       // lSm should be: {-1, 0, -1, -1, 3, 3}
    
31.       for (int i = 0; i < n; ++i) {
    
32.           while (stk.size() && (nums[stk.top()] > nums[i] || nums[stk.top()] == nums[i] - 1)) stk.pop();
    
33.           if (stk.size()) lSm[i] = stk.top();
    
34.           stk.push(i);
    
35.       }
    
36.       // rSm should be: {-1, 3, 3, -1, -1}
    
37.       while (stk.size()) stk.pop();
    
38.       for (int i = n - 1; i >= 0; --i) {
    
39.           while (stk.size() && (nums[stk.top()] > nums[i] || nums[stk.top()] == nums[i] - 1)) stk.pop();
    
40.           if (stk.size()) rSm[i] = stk.top();
    
41.           stk.push(i);
    
42.       }
    
43.       long long ttl = 0;
    
44.       for (int i = 0; i < n; ++i) {
    
45.           int rChoice = 0, lChoice = 0;
    
46.           long long subttl = 0;
    
47.           if (idxM1[i] < i) {
    
48.               if (rSm[i] != -1) {
    
49.                   // ... 3 4 1 ... [n=7]
    
50.                   rChoice = n - rSm[i];
    
51.                   lChoice = i - idxM1[i];
    
52.                   subttl += lChoice * rChoice;
    
53.               }
    
54.               if (lSm[i] > idxM1[i]) {
    
55.                   // ... 3 1 4 ... [n=7]
    
56.                   lChoice = lSm[i] - idxM1[i];
    
57.                   rChoice = n - i;
    
58.                   subttl += lChoice * rChoice;
    
59.               }
    
60.               if (rSm[i] != -1 && lSm[i] > idxM1[i]) subttl--;
    
61.           }
    
62.           else {  // idxM1[i] > i
    
63.               if (lSm[i] != -1) {
    
64.                   // ... 1 4 3 ... [n=7]
    
65.                   rChoice = idxM1[i] - i;
    
66.                   lChoice = lSm[i] + 1;
    
67.                   subttl += lChoice * rChoice;
    
68.               }
    
69.               if (rSm[i] != -1 && rSm[i] < idxM1[i]) {
    
70.                   // ... 1 4 3 ... [n=7]
    
71.                   lChoice = i + 1;
    
72.                   rChoice = idxM1[i] - rSm[i];
    
73.                   subttl += lChoice * rChoice;
    
74.               }
    
75.               if (lSm[i] != -1 && rSm[i] != -1 && rSm[i] < idxM1[i]) subttl--;
    
76.           }
    
77.           cout << nums[i] << ":" << (subttl > 0 ? subttl : 0) << endl;
    
78.           ttl += (subttl > 0 ? subttl : 0);
    
79.       }
    
80.       return ttl;
    
81.   }

复制代码

补充内容 (2022-09-30 14:38 +8:00):
前半部分都是一樣的。從那個算ttl的地方開始有不同。這個調試版本能夠打印出來每一個數字對ttl的貢獻，可能如果還有bug的話找起來會方便一點。

bydstand 发表于 2022-9-29 16:51
是的，根本不是单调栈

https://leetcode.com/playground/55hXJqEY

用單調棧才能夠線性時間找到lSm和rSm兩個數組，所以是讓整個解法只需線性時間的關鍵。

求米，很缺米🤣
花了3个小时才弄懂这个题目

LC2200的飘过 不知道思路对不对哈

每个数字m能否贡献一个imbal 取决于他的直系学长m+1在不在这个子数组里，那么我们loop一遍1到n-1，去找他的直系学长的位置，可以用map辅助查找，比如对于数字3(index 2), 数字4位于index 6，那么数字3对于最终结果的贡献就是包含index2且不包含index6的子数组的个数，也就是（2-0+1）* （6-2），这些贡献值累加即可

我的solution, both time and space complexity O(N).
抱歉用的java，觉得可读性比python高点
https://leetcode.com/playground/c2dq4ULy

PerryLu 发表于 2022-9-30 15:38
求米，很缺米🤣
花了3个小时才弄懂这个题目

老哥，这个是改写heap的比较规则，
this.maxHeap = new PriorityQueue<>((a, b) -> {
                        if(b.getValue().equals(a.getValue()))
                                return b.getKey().compareTo(a.getKey());
                        return b.getValue().compareTo(a.getValue());       
                });

这个代码用Python写了以后是下面这样吗？谢谢老哥！！！
class MaxHeapItem:
    def __init__(self, name, score):
        self.name = name
        self.score = score
    def __lt__(self, other):
        return self.score > other.score or \
               (self.score == other.score and self.name < other.name)