日渐头秃的代码日记 -- 第294场周赛第4题

113 阅读2分钟

巫师的总力量和

作为国王的统治者,你有一支巫师军队听你指挥。

给你一个下标从 0 开始的整数数组 strength ,其中 strength[i] 表示第 i 位巫师的力量值。对于连续的一组巫师(也就是这些巫师的力量值是 strength 的 子数组),总力量 定义为以下两个值的 乘积 :

  • 巫师中 最弱 的能力值。
  • 组中所有巫师的个人力量值 之和 。

请你返回 所有 巫师组的  力量之和。由于答案可能很大,请将答案对 109 + 7 取余 后返回。

子数组 是一个数组里 非空 连续子序列。

 

示例 1:

输入: strength = [1,3,1,2]
输出: 44
解释: 以下是所有连续巫师组:
- [1,3,1,2][1] ,总力量值为 min([1]) * sum([1]) = 1 * 1 = 1
- [1,3,1,2][3] ,总力量值为 min([3]) * sum([3]) = 3 * 3 = 9
- [1,3,1,2][1] ,总力量值为 min([1]) * sum([1]) = 1 * 1 = 1
- [1,3,1,2][2] ,总力量值为 min([2]) * sum([2]) = 2 * 2 = 4
- [1,3,1,2][1,3] ,总力量值为 min([1,3]) * sum([1,3]) = 1 * 4 = 4
- [1,3,1,2][3,1] ,总力量值为 min([3,1]) * sum([3,1]) = 1 * 4 = 4
- [1,3,1,2][1,2] ,总力量值为 min([1,2]) * sum([1,2]) = 1 * 3 = 3
- [1,3,1,2][1,3,1] ,总力量值为 min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5
- [1,3,1,2][3,1,2] ,总力量值为 min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6
- [1,3,1,2][1,3,1,2] ,总力量值为 min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7
所有力量值之和为 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44

示例 2:

输入: strength = [5,4,6]
输出: 213
解释: 以下是所有连续巫师组:
- [5,4,6][5] ,总力量值为 min([5]) * sum([5]) = 5 * 5 = 25
- [5,4,6][4] ,总力量值为 min([4]) * sum([4]) = 4 * 4 = 16
- [5,4,6][6] ,总力量值为 min([6]) * sum([6]) = 6 * 6 = 36
- [5,4,6][5,4] ,总力量值为 min([5,4]) * sum([5,4]) = 4 * 9 = 36
- [5,4,6][4,6] ,总力量值为 min([4,6]) * sum([4,6]) = 4 * 10 = 40
- [5,4,6][5,4,6] ,总力量值为 min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60
所有力量值之和为 25 + 16 + 36 + 36 + 40 + 60 = 213

 

提示:

  • 1 <= strength.length <= 10^5
  • 1 <= strength[i] <= 10^9

题目解析

周赛时候尝试过的用户有1600+人,仅有140+人通过了,可见本题难度之大。

第一个方法是直接暴力求解,即从第i个位置到第j个位置,计算一次总力量值,最后把所有的力量值加起来,就可以得到最后的结果,这种方法可以通过67个测试case,一共80个case还有一部分无法通过,可见必须要进行优化。

每一个数字,都有可能成为某一个子序列中最小的那个数,因此只需要把每一个数能带来的所有贡献力量值加到一起,就可以得到总的力量值。

以[2, 5, 4, 3, 4, 2]这样一串数字为例,假设当前要计算的是数字3的贡献值,它能带来贡献的子序列,只有可能是在 [5, 4, 3, 4] 这样一个边界范围以内且包含它自己,因此就只有 [3], [4, 3], [5, 4, 3], [3, 4], [4, 3, 4], [5, 4, 3, 4]一共6种情况。 假设目前有前缀和S,对应原始序列的S为 [2, 7, 9, 12, 16, 18],那么数字3贡献的范围一共是

contribute = S3-S2 + S3-S1 + S3-S0 + S4-S2 + S4-S1 + S4-S0
= 3*(S4+S3) - 2*(S2+S1+S0)

因为要计算S2+S1+S0,相当于再算一遍前缀和,因此可以对S进行前缀和计算得到原始序列的二次前缀和。 只需要找到每一个数字成为最小数字的区间范围即可,即左边最近的一个比自己小的作为起点,右边第一个不比自己大的作为终点(不能两边都是严格小于自己,那样会有子序列遗漏)

class Solution:
    def totalStrength(self, strength: List[int]) -> int:
        MOD = 10 ** 9 + 7

        n = len(strength)
        left = [-1] * n
        stack = []
        for i, v in enumerate(strength):
            while stack and strength[st[-1]] >= v:
                stack.pop()
            if stack:
                left[i] = st[-1]
            stack.append(i)

        right= [n] * n
        stack = []
        for i in range(n - 1, -1, -1):
            while stack and strength[st[-1]] > strength[i]:
                stack.pop()
            if stack:
                right[i] = st[-1]
            stack.append(i)


        s = [0] * (n + 1)
        for i, v in enumerate(strength):
            s[i + 1] = (s[i] + v) % MOD
        ss = [0] * (n + 2)
        for i, v in enumerate(s):
            ss[i + 1] = (ss[i] + v) % MOD

        ans = 0
        for i, v in enumerate(strength):
            l = left[i] + 1
            r = right[i] - 1
            total = ((i - l + 1) * (ss[r + 2] - ss[i + 1]) - (r - i + 1) * (ss[i + 1] - ss[l])) % MOD
            ans = (ans + total * v) % MOD
        return ans