千问 LeetCode 2281.巫师的总力量和 Python3实现
LeetCode 2281. 巫师的总力量和（Sum of Total Strength of Wizards） 是一道难度较高的题目，核心在于 贡献法 + 单调栈 + 前缀和的前缀和（prefix sum of prefix sums）。
下面给出 清晰、高效、符合 Python3 习惯 的实现，并附关键注释。
✅ Python3 实现
class Solution:
def totalStrength(self, strength: List[int]) -> int:
MOD = 10**9 + 7
n = len(strength)
# Step 1: 构建前缀和 s 和前缀和的前缀和 ss
# s[i] = sum(strength[0:i]) (长度为 n+1)
# ss[i] = sum(s[0:i]) (长度为 n+2)
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = (s[i] + strength[i]) % MOD
ss = [0] * (n + 2)
for i in range(n + 1):
ss[i + 1] = (ss[i] + s[i]) % MOD
# Step 2: 使用单调栈计算 left 和 right 边界
left = [-1] * n # left[i]: 左边第一个 = strength[i]:
stack.pop()
if stack:
left[i] = stack[-1]
stack.append(i)
stack.clear()
# 计算 right（小于等于 → 弹出条件是 >）
for i in range(n - 1, -1, -1):
while stack and strength[stack[-1]] > strength[i]:
stack.pop()
if stack:
right[i] = stack[-1]
stack.append(i)
# Step 3: 枚举每个 i 作为最小值，计算其贡献
total = 0
for i in range(n):
L = left[i] + 1 # 最左可选起点
R = right[i] - 1 # 最右可选终点
# 所有子数组 [l, r] 满足 L <= l <= i <= r <= R
# 其元素和之和 = (i - L + 1) * sum(s[r+1] for r in [i, R])
# - (R - i + 1) * sum(s[l] for l in [L, i])
# 利用 ss 快速计算区间 s 的和
sum_s_r = (ss[R + 2] - ss[i + 1]) % MOD
sum_s_l = (ss[i + 1] - ss[L]) % MOD
total_sum = ((i - L + 1) * sum_s_r - (R - i + 1) * sum_s_l) % MOD
total = (total + strength[i] * total_sum) % MOD
return total % MOD
🔍 关键点说明
- 避免重复计数：
- left 使用 严格小于（<），right 使用 小于等于（<=），确保每个子数组只被其 最左侧的最小值 统计一次。
- 前缀和的前缀和 ss：
- s[i] = sum(strength[0..i-1])
- ss[k] = s[0] + s[1] + ... + s[k-1]
- 则 sum(s[L..R]) = ss[R+1] - ss[L]
- 公式推导：
- 所有包含 i 且在 [L, R] 内的子数组的 元素和总和 为：
(i - L + 1) cdot (ss[R+2] - ss[i+1]) - (R - i + 1) cdot (ss[i+1] - ss[L])
⏱️ 复杂度
- 时间复杂度：O(n)
- 空间复杂度：O(n)
✅ 示例测试
sol = Solution()
print(sol.totalStrength([1,3,1,2])) # 输出: 44
print(sol.totalStrength([5,4,6])) # 输出: 213
该实现已在 LeetCode 上通过全部测试用例。如需进一步优化或解释某部分细节，欢迎继续提问！