🔗 🔴 2281. Sum of Total Strength of Wizards

Problem Statement

題目簡述

給定一個長度為 nn 的整數陣列 strength，其中 strength[i] 表示第 ii 個巫師的力量值。
對於任意一個連續的巫師子陣列，其「總力量值」定義為：該子陣列中最弱巫師的力量值 ×\times 該子陣列中所有巫師力量值的總和。
請返回所有連續子陣列的總力量值之和。由於答案可能很大，請將結果對 109+710^9 + 7 取模後返回。

Constraints

約束條件

  • 1≤strength.length≤1051 \le \text{strength.length} \le 10^5
  • 1≤strength[i]≤1091 \le \text{strength}[i] \le 10^9

思路：貢獻法 + 單調堆疊 + 前綴和的前綴和

1. 尋找突破口：從「子陣列」到「單一元素的貢獻」

如果直接枚舉所有子陣列，再求每個子陣列的最小值與區間和，時間複雜度至少是 O(n2)\mathcal{O}(n^2)，對於 n≤105n \le 10^5 的數據範圍顯然會逾時。

核心套路：貢獻法

當我們需要對「所有子陣列的某種統計量」求和時，不妨轉換視角：
不要去枚舉子陣列，而是去枚舉每個元素，看它作為「最小值」時，能對哪些子陣列做出貢獻。

對於陣列中的某個元素 xx（設其下標為 ii）：

  • 假設以 xx 為最小值的子陣列，其左端點最左能延伸到 ll，右端點最右能延伸到 rr。
  • 那麼，任何左端點在 [l,i][l, i]、右端點在 [i,r][i, r] 之間的子陣列，其最小值都是 xx。
  • 這些子陣列的左端點有 a=i−l+1a = i - l + 1 種選擇，右端點有 b=r−i+1b = r - i + 1 種選擇。

2. 確定邊界：單調堆疊與重複元素的處理

如何高效找到每個元素作為最小值的左右邊界？這可以使用單調堆疊（Monotonic Stack）來解決。

以求右邊界為例，從左到右掃描並維護一個單調遞增的堆疊。當目前元素比堆疊頂端更小時，表示堆疊頂端對應的位置第一次在右側遇到更小元素，因此可以確定它的右邊界就是目前位置。左邊界同理，只是改成從右往左掃描，並配合重複元素的歸屬規則調整比較條件。

易錯點：重複元素的邊界定義

如果陣列中存在相同的元素（例如 [1, 2, 1]），在定義邊界時如果不加小心，同一個子陣列可能會被多個相同的最小值重複計算。

解決方案：
我們可以人為規定，對於多個相同的最小值，只在其中某一個位置（例如最右側或最左側）計算貢獻。
在實作上，我們可以定義：

  • 左邊界：左側第一個小於或等於目前元素的位置。
  • 右邊界：右側第一個嚴格小於目前元素的位置。

這樣一來，每個子陣列的最小值就有了唯一的歸屬，既不重複也不遺漏。

除了上述做法外，還有其他的做法，甚至可以做到在 1-pass 中同時求出左右邊界並計算貢獻，但為了清晰起見，我們這裡分兩次掃描。

3. 核心推導：前綴和的前綴和

確定了左端點範圍 [l,i][l, i] 與右端點範圍 [i,r][i, r] 後，我們需要計算所有這些子陣列的元素和之和。
設前綴和陣列為 SS，其中 S[k]=∑j=0k−1strength[j]S[k] = \sum_{j=0}^{k-1} \text{strength}[j]。那麼任意子陣列 [left,right][\textit{left}, \textit{right}] 的元素和可以表示為 S[right+1]−S[left]S[\textit{right}+1] - S[\textit{left}]。

我們要計算的總和為：

∑j=li∑k=ir(S[k+1]−S[j])\sum_{j=l}^{i} \sum_{k=i}^{r} (S[k+1] - S[j])

將這個雙重求和展開：

∑j=li∑k=ir(S[k+1]−S[j])=∑j=li∑k=irS[k+1]−∑j=li∑k=irS[j]=(i−l+1)∑k=irS[k+1]−(r−i+1)∑j=liS[j]\begin{aligned} \sum_{j=l}^{i} \sum_{k=i}^{r} (S[k+1] - S[j]) &= \sum_{j=l}^{i} \sum_{k=i}^{r} S[k+1] - \sum_{j=l}^{i} \sum_{k=i}^{r} S[j] \\ &= (i - l + 1) \sum_{k=i}^{r} S[k+1] - (r - i + 1) \sum_{j=l}^{i} S[j] \end{aligned}

為了在 O(1)\mathcal{O}(1) 時間內算出 ∑S[k+1]\sum S[k+1] 和 ∑S[j]\sum S[j]，我們可以再對前綴和陣列 SS 再做一次前綴和，即前綴和的前綴和（設為 SSSS），其中 SS[k]=∑j=0k−1S[j]SS[k] = \sum_{j=0}^{k-1} S[j]。

利用 SSSS，我們可以將上述兩部分區間和轉化為：

  • ∑k=irS[k+1]=SS[r+2]−SS[i+1]\sum_{k=i}^{r} S[k+1] = SS[r+2] - SS[i+1]
  • ∑j=liS[j]=SS[i+1]−SS[l]\sum_{j=l}^{i} S[j] = SS[i+1] - SS[l]

帶回原式，每個元素作為最小值的子陣列元素和之和為：

(i−l+1)⋅(SS[r+2]−SS[i+1])−(r−i+1)⋅(SS[i+1]−SS[l])(i - l + 1) \cdot (SS[r+2] - SS[i+1]) - (r - i + 1) \cdot (SS[i+1] - SS[l])

最後，將此結果乘以目前元素的值 xx，並對 109+710^9 + 7 取模，累加到答案中即可。

複雜度分析

  • 時間複雜度：O(n)\mathcal{O}(n)，其中 nn 為陣列的長度。
  • 空間複雜度：O(n)\mathcal{O}(n)。

類題

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
MOD = int(1e9) + 7

class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        n = len(strength)

        # L[i] 表示左側第一個 <= strength[i] 的位置
        # R[i] 表示右側第一個 < strength[i] 的最小位置
        L = [-1] * n
        R = [n] * n
        st = []
        for i, x in enumerate(strength):
            while st and x < strength[st[-1]]:
                R[st.pop()] = i
            st.append(i)

        st = []
        for i in range(n - 1, -1, -1):
            x = strength[i]
            while st and x <= strength[st[-1]]:
                L[st.pop()] = i
            st.append(i)

        s = list(accumulate(strength, initial=0))
        ss = list(accumulate(s, initial=0))

        ans = 0
        for i, x in enumerate(strength):
            # 以 strength[i] 作為最小值的區間左端點可以為 [L[i] + 1, i]，右端點可以為 [i, R[i] - 1]
            l, r = L[i] + 1, R[i] - 1
            a, b = i - l + 1, r - i + 1
            tot = a * (ss[r + 2] - ss[i + 1]) - b * (ss[i + 1] - ss[l])
            ans += tot * x % MOD
            ans %= MOD
        return ans