大厂真题：【前缀和】米哈游2023秋招-米小游的极差之和

题目描述与示例

题目描述

米小游拿到了一个数组a，她用这个数组构造一个新数组b，其中ai代表b数组中有ai个i。

例如，若a = [2,3,1]，那么b = [1,1,2,2,2,3]，因为a1=2，代表b 数组中有2个1；a2=3，代表b数组中有3个2；a3 = 1，代表b数组中有1个3。

现在给定a数组，你需要帮米小游求出b数组中所有连续子数组的极差之和。由于答案可能过大，请对10^9+7 取模。

数组的极差指最大值减去最小值。

输入描述

第一行输入一个正整数n，代表a数组的元素数量。

第二行输入n 个正整数ai，代表a数组的元素。

1 ≤ n ≤ 10^5
1 ≤ ai ≤ 10^9

输出描述

一个整数，代表数组中所有区间的极差之和，对10^9+7取模的值。

示例

输入

Plain 复制代码

2
2 1

输出

Plain 复制代码

说明

a=[2,1]时，b数组为[1,1,2]。

此时b数组共有6个连续子数组：

[1]的极差为0。

[2]的极差为0。

[1,1]的极差为0。

[1,2]的极差为1。

[1,1,2]的极差为1。

因此答案是0+0+0+0+1+1=2。

解题思路

根据数据范围大小，显然不能将a数组展开回b数组求解，而应该直接从a数组出发解决问题。

注意到b数组实际上是一个单调非递减数组，其子数组也一定单调非递减数组。当子数组的最大值和最小值不相等（即最后一个元素和第一个元素）时，该子数组才对极差有贡献。

对于任意的两元组 ( i , j ) (i, j) (i,j)满足 1 < = i < j < = n 1 <= i < j <= n 1<=i<j<=n，a数组中均有ai个i和aj个j，根据乘法原理 ，这对二元组对总极差的贡献为 ∑ j = i + 1 n ( j − i ) a j a i \sum_{j=i+1}^{n}(j-i)a_ja_i j=i+1∑n(j−i)ajai。考虑所有的二元组，则总极差为为 ∑ i = 1 n ∑ j = i + 1 n ( j − i ) a j a i \sum_{i=1}^{n}\sum_{j=i+1}^{n}(j-i)a_ja_i i=1∑nj=i+1∑n(j−i)ajai

如果直接对上式进行求和计算，需要使用双重循环，时间复杂度为O(n^2)，会出现部分用例超时的情况。故需要对上述式子进行恒等变换，再观察是否有降低复杂度的做法。

∑ i = 1 n ∑ j = i + 1 n ( j − i ) a j a i = ∑ i = 1 n a i ⋅ ∑ j = i + 1 n ( j − i ) a j = ∑ i = 1 n a i ⋅ ∑ j = i + 1 n ( j a j − i a j ) = ∑ i = 1 n a i ⋅ ( ∑ j = i + 1 n j a j − ∑ j = i + 1 n i a j ) \sum_{i=1}^{n}\sum_{j=i+1}^{n}(j-i)a_ja_i = \\ \sum_{i=1}^{n}a_i ·\sum_{j=i+1}^{n}(j-i)a_j = \\ \sum_{i=1}^{n}a_i ·\sum_{j=i+1}^{n}(ja_j-ia_j) = \\ \sum_{i=1}^{n}a_i ·(\sum_{j=i+1}^{n}ja_j-\sum_{j=i+1}^{n}ia_j) i=1∑nj=i+1∑n(j−i)ajai=i=1∑nai⋅j=i+1∑n(j−i)aj=i=1∑nai⋅j=i+1∑n(jaj−iaj)=i=1∑nai⋅(j=i+1∑njaj−j=i+1∑niaj)

在最终的结果中，可以看出对于每一个i，其极差的情况都仅和j = i+1相关，对于式子中的两个与j相关的求和 ∑ j = i + 1 n j a j \sum_{j=i+1}^{n}ja_j j=i+1∑njaj和 ∑ j = i + 1 n a j \sum_{j=i+1}^{n}a_j j=i+1∑naj，我们可以用类似构建前缀和的方式，

令 P i = ∑ j = 1 i j a j P_i = \sum_{j=1}^{i}ja_j Pi=j=1∑ijaj，那么存在 ∑ j = i + 1 n j a j = ∑ j = 1 n j a j − ∑ j = 1 i j a j = P n − P i \sum_{j=i+1}^{n}ja_j = \sum_{j=1}^{n}ja_j - \sum_{j=1}^{i}ja_j = P_n-P_i j=i+1∑njaj=j=1∑njaj−j=1∑ijaj=Pn−Pi

令 Q i = ∑ j = 1 i a j Q_i = \sum_{j=1}^{i}a_j Qi=j=1∑iaj，那么存在 ∑ j = i + 1 n a j = ∑ j = 1 n a j − ∑ j = 1 i a j = Q n − Q i \sum_{j=i+1}^{n}a_j = \sum_{j=1}^{n}a_j - \sum_{j=1}^{i}a_j = Q_n-Q_i j=i+1∑naj=j=1∑naj−j=1∑iaj=Qn−Qi

在O(1)的时间复杂度下求得这两个式子的结果了。（PS：严谨来说，这应该叫做后缀和才比较合适

代码

Python

Python 复制代码

# 题目：【前缀和】米哈游2023秋招-米小游的极差之和
# 作者：闭着眼睛学数理化
# 算法：前缀和
# 代码有看不懂的地方请直接在群上提问


from itertools import accumulate

n = int(input())
a_list = list(map(int, input().split()))

MOD = 10**9+7
# 分别构建p_list和q_list
p_list = [0] + list(accumulate(i*ai for i, ai in enumerate(a_list, 1)))
q_list = [0] + list(accumulate(ai for ai in a_list))
# 由于p_list和q_list的最后一个数需要反复用到，所以提前取出
pn = p_list[-1]
qn = q_list[-1]

ans = 0
# 遍历1到n的所有i，根据恒等式转换得到的结果进行计算
for i in range(1, n+1):
    # 注意这里i是从1开始，a_list中是从索引0开始的
    # 故取ai时，对应的是a_list[i-1]
    ans += a_list[i-1]*(pn-p_list[i]-(qn-q_list[i])*i)
    ans %= MOD

print(ans)

Java

Java 复制代码

import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        long[] aList = new long[n];

        final long MOD = 1000000007;
        long ans = 0;

        for (int i = 0; i < n; i++) {
            aList[i] = sc.nextLong();
        }

        long[] pList = new long[n + 1];
        long[] qList = new long[n + 1];

        for (int i = 1; i <= n; i++) {
            pList[i] = pList[i - 1] + (long) i * aList[i - 1];
            qList[i] = qList[i - 1] + aList[i - 1];
        }

        long pn = pList[n];
        long qn = qList[n];

        for (int i = 1; i <= n; i++) {
            ans += aList[i - 1] * (pn - pList[i] - (qn - qList[i]) * i);
            ans %= MOD;
        }

        System.out.println(ans);
    }
}

C++

C++ 复制代码

#include <iostream>
#include <vector>
#include <numeric>

using namespace std;

int main() {
    int n;
    cin >> n;
    vector<int> a_list(n);

    const int MOD = 1000000007;
    long long ans = 0;

    for (int i = 0; i < n; i++) {
        cin >> a_list[i];
    }

    vector<long long> p_list(n + 1, 0);
    vector<long long> q_list(n + 1, 0);

    for (int i = 1; i <= n; i++) {
        p_list[i] = p_list[i - 1] + static_cast<long long>(i) * a_list[i - 1];
        q_list[i] = q_list[i - 1] + a_list[i - 1];
    }

    long long pn = p_list[n];
    long long qn = q_list[n];

    for (int i = 1; i <= n; i++) {
        ans += static_cast<long long>(a_list[i - 1]) * (pn - p_list[i] - (qn - q_list[i]) * i);
        ans %= MOD;
    }

    cout << ans << endl;
    
    return 0;
}

时空复杂度

时间复杂度：O(N)。化简后，仅需一次遍历数组。

空间复杂度：O(N)。两个后缀和数组所占用空间。

华为OD算法/大厂面试高频题算法练习冲刺训练

华为OD算法/大厂面试高频题算法冲刺训练目前开始常态化报名！目前已服务100+同学成功上岸！
课程讲师为全网50w+粉丝编程博主@吴师兄学算法 以及小红书头部编程博主@闭着眼睛学数理化
每期人数维持在20人内，保证能够最大限度地满足到每一个同学的需求，达到和1v1同样的学习效果！
60+天陪伴式学习，40+直播课时，300+动画图解视频，300+LeetCode经典题，200+华为OD真题/大厂真题，还有简历修改、模拟面试、专属HR对接将为你解锁
可上全网独家的欧弟OJ系统练习华子OD、大厂真题
可查看链接 OD算法冲刺训练课程表 & OD真题汇总(持续更新)
绿色聊天软件戳 od1336了解更多