计算元素贡献 专题

本专题尝试用比较容易理解的方式，去解释一系列计算元素贡献的一类题的通用解法与思考方式

序言

2022.5.22 上午的周赛T4 6077. 巫师的总力量和 - 力扣（LeetCode）就是这一类典型的题型，对于这类 统计序列中所有子数组的某个性质/值的总和 的问法，暴力枚举是一个显然的做法，但是通常它的复杂度会达到$O(n^2)$或者$O(n^3)$的级别，如果题目数据量在$10^4$以上，绝对会超时。那么，对于这一类问题，我们关注如何用简单的思路去实现一个$O(n)$的解法。

这个思路，用一句话来说就是：统计序列中 每个数/每种长度 产生的贡献。我们遍历序列，对于当前元素，如果我们能在$O(1)$的时间内计算出当前 元素/长度 对于答案产生的贡献值，我们就能够将时间复杂度控制在$O(n)$。

对于这样一个方法论，实际上有很多种类型的题目都可以应用，接下来由浅入深，结合例题进行解释。

较直接的问题

问题1：求所有子数组的和

抛砖引玉，先从最直接的一道题出发：1814 · 所有子数组之和 - LintCode 题目的意思是：

给出一个长度为 $n$ 的序列$[a_1,a_2,…,a_n]$（如无特殊说明，此后所有序列定义都由下标1开始），求这个序列的所有子数组的和，例如，序列$[1,2,3]$的所有子数组为$[1],[2],[3],[1,2],[2,3],[1,2,3]$，它们的和为$1+2+3+1+2+2+3+1+2+3=20$。

乍一看这题，有一个显然的解法，就是我们可以枚举每一个子数组，即 一层循环枚举左端点，一层枚举右端点，然后暴力求这个子数组的和（这个可以通过前缀和优化成$O(1)$求），因此复杂度为$O(n^2)$。

如果用统计贡献的思路，那么对于这样一个问题，通常有两种思考方式：

  1. 考虑计算每个元素对答案的贡献
  2. 采用动态规划的思路，考虑在遍历这个数组的时候，加入一个新的数对答案的贡献

思路1. 考虑计算每个元素对答案的贡献

我们先思考第1种思路（也是本问题最简单的思路），举个栗子，观察这样一个长度为4的序列$[1,3,5,2]$，思考对于$3$这个元素，它会出现在哪些子数组中？答案是$6$个子数组：$[3],[3,5],[3,5,2],[1,3],[1,3,5],[1,3,5,2]$，因此，它会对答案产生$6$次贡献，则它的贡献为$3\times 6=18$。

进一步地，仔细观察这些子数组的左右端点，我们能发现这样一个规律：左端点一定是小于等于当前数的位置，右端点一定大于等于当前数的位置（因为必须包含这个数），我们发现，$3$这个数处在的位置为$2$，因此它出现在$2\times (4-2+1)=6$个子数组中。更一般地，利用乘法原理我们知道，对于一个下标为$i$的数$a[i]$，它一定会出现在$i\times(n-i+1)$个子数组中（即 左侧选$i$个，右侧选$n-i+1$个，乘法原理），因此它对答案的贡献为$a[i]\times i\times (n-i+1)$。

遍历整个数组，对每个元素直接$O(1)$求出这个值累加到答案中，就完成了$O(n)$的解法。

1
2
3
int ans = 0;
for(int i=1;i<=n;++i)
    ans += a[i] * i * (n-i+1);

这里我们来模拟一下计算流程，以 $[1,3,5,2]$为例。

  • 考虑$a_1 = 1$计算对答案的贡献为$1\times 1\times 4 = 4$，加入至答案 $ans = 4$
  • 考虑$a_2 = 3$计算对答案的贡献为$3\times 2\times 3 = 18$，加入至答案 $ans=4+18=22$
  • 考虑$a_3 = 5$计算对答案的贡献为$5\times 3\times 2 = 30$，加入至答案 $ans=22+30=52$
  • 考虑$a_4 = 2$计算对答案的贡献为$2\times 4\times 1 = 8$，加入至答案 $ans=52+8=60$

思路2. 动态规划，考虑加入一个新数对答案的贡献

接下来考虑动态规划的思路，我们定义一个状态$f(i)$表示处理完前$i$个数后的答案值。那么，对于$f(i-1)$已经求出的情况下，考虑新加入元素$a[i]$对答案的贡献。

此时，我们还是需要考虑元素$a[i]$的加入相比于前$i-1$个元素产生的贡献，多出了哪些贡献，这个贡献实际上就是 多出来的子数组的和，进一步地，我们发现这些多出来的子数组实际上都是以$a[i]$为结尾的子数组，而这些子数组，显然可以把它拼到以$a[i-1]$为结尾的子数组后面，这样以来，就相当于多出了$(i-1)$个（包含了$a[i]$的）子数组，此外，$a[i]$本身也可以自己作为一个新的子数组，所以总共产生了$(i-1+1=i)$个新的子数组，因此这里产生了一个值为$(i\times a[i])$的贡献，这里要特别注意，这个贡献实际上代表的是：添加$a[i]$后，相较于以$a[i-1]$为结尾的所有子数组的和 多出来的部分，因此我们需要再维护一个状态$g(i)$代表以$a[i]$为结尾的所有子数组的和，那么就有状态转移方程：$g(i) = g(i-1) + i\times a[i]$。对于$f(i)$的状态转移，就是把$g(i)$这部分加上就行了：$f(i) = f(i-1) + g(i)$。

更进一步思考，实际上$g(i)$维护了一个类似前缀和的东西：$\sum_\limits{i=1}^n i\times a[i]$

代码实现方面，由于状态$i$仅与$i-1$有关，因此可以滚动数组将空间优化为$O(1)$。

1
2
3
4
5
6
int f = 0; // 答案
int g = 0; // 以a_i为结尾的所有子数组的和
for(int i=1;i<=n;++i){
    g += i * a[i];
    f += g;
}

这里我们来模拟一下计算流程，还是刚刚的例子 $[1,3,5,2]$。

  • 考虑第1个数$1$，以$a_1 = 1$为结尾的所有子数组和为：$1\times 1 = 1$（即$[1]$这一个子数组），此后计算对答案的贡献，$f = 1$；

  • 考虑第2个数$3$，以$a_2 = 3$为结尾的所有子数组和为：$1+2\times 3 = 7$（即$[1,3],[3]$这两个子数组），此后计算对答案的贡献，$f = 1+7 = 8$；

  • 考虑第3个数$5$，以$a_3 = 5$为结尾的所有子数组和为：$7+3\times 5 = 22$（即$[1,3,5],[3,5],[5]$这三个子数组），此后计算对答案的贡献，$f = 8+22 =30$；

  • 考虑第4个数$2$，以$a_4 = 2$为结尾的所有子数组和为：$22+4\times 2 = 30$（即$[1,3,5,2],[3,5,2],[5,2],[2]$这四个子数组），此后计算对答案的贡献，$f = 30+30 = 60$

因此最终答案为$f=60$。

这样一个直接的问题，我们从两种完全不同的思路出发，得到了相同的结果。因此，对于此类计算贡献的问题，我们通常从这两种思路出发去寻找问题的解法，当然，其他很多问题并没有此题这么直接，通常情况下，他们会结合 数学、单调栈/单调队列、动态规划 等知识点综合考察。

问题1变式：求序列所有子集和

如果将问题1种的 子数组 改为 子集 ，则如何求解呢？

这个问题并不难，可以从思路1的角度出发，还是乘法原理，但左侧端点和右侧端点的选择变多了，对于$a[i]$来说，左侧能够选$2^{i-1}$个子集，右侧能够选$2^{n-i}$个子集，因此对答案的贡献为$a[i] \times 2^{i-1}\times 2^{n-i}$.

1
2
3
int ans = 0;
for(int i=1;i<=n;++i)
    ans += a[i] * (1<<i-1) * (1<<n-i);

问题1进阶：求所有子数组平均数的和

问题1的进阶版本：

给出一个长度为 $n$ 的序列$[a_1,a_2,…,a_n]$，求这个序列的所有子数组的平均数，例如，序列$[1,3,5]$的所有子数组为$[1],[3],[5],[1,3],[3,5],[1,3,5]$，它们的平均数的和为$1+3+5+\frac{1+3}{2}+\frac{3+5}{2}+\frac{1+3+5}{3}=18$。

由于最终答案可能出现浮点数，因此最终输出结果有两种方式：

  1. 输出浮点数
  2. 输出模$10^9+7$意义下的答案

我们先考虑输出浮点数的情况，对于这样一个问题，我们同样考虑如何计算每个元素产生的贡献。我们发现，对于不同的子数组的长度，每个元素产生的贡献的权重是不一样的，比如说对于$[1,3,5]$这个序列来说，元素$1,3,5$对答案产生的贡献如下：

似乎可以看出有一种规律，我们把元素数写多一些，例如$[a_1,a_2,a_3,a_4,a_5]$种每个元素对答案的贡献分别为：

可以看出，右侧的值实际上是序列的$1+\frac12 + \frac13+…+\frac1n$某种加权和，而这个权重的规律符合三角波的规律，即：

对于$a1$来说，就是$\sum \limits_{i=1}^{n} \frac{1}{i}$；

对于$a2$来说，就是加上了$\sum \limits_{i=2}^{n-1} \frac{1}{i}$；

对于$a3$来说，就是再加上$\sum\limits_{i=3}^{n-2} \frac{1}{i}$，以此类推。

而对于过半之后的元素，也是经过类似的操作，只不过变成了减法。

因此，我们可以通过$O(n)$预处理出$\sum\limits_{i=1}^n \frac{1}{i}$这样一个前缀和，然后在遍历序列的时候$O(1)$处理变化值就行了，总复杂度$O(n)$。

对于输出模$10^9+7$意义下的答案这样的要求，只需$O(n)$预处理出$[1,2,,…,n]$这些元素的乘法逆元即可，复杂度依然是$O(n)$。

解法来源于我在StackExchange上问的相同的问题：algorithms - Sum of average of all subarrays - Computer Science Stack Exchange

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
#include<bits/stdc++.h>
using namespace std;
constexpr int mod = 1e9+7;
void solve(){
    int n;
    cin>>n;
    vector<int> a(n+1);
    vector<int> inv(n+1);
    for(int i=1;i<=n;++i) cin>>a[i];
    inv[1] = 1;
    for(int i=2;i<=n;++i) // linearly calculate mod inverse of [1,2,...,n]
        inv[i] = 1LL * (mod-mod/i)*inv[mod%i]%mod;
    for(int i=2;i<=n;++i) // prefix of mod inverse of [1,2,...,n]
        inv[i] = (1LL * inv[i] + inv[i-1])%mod;
    int ans = 0;
    int tot = 0;
    int r = n, l = 0;
    for(int i=1;i<=(n+1)/2;++i){
        tot = (1LL * tot + inv[r--] - inv[l++] + mod) % mod;
        ans = (1LL * ans + 1LL * a[i] * tot ) % mod;
    }
    if(n%2==0) ans = (1LL * ans + 1LL * a[n/2+1] * tot) % mod;
    for(int i=(n+1)/2+1+(n%2==0);i<=n;++i){
        tot = (1LL * tot - inv[++r] + inv[--l] + mod) % mod;
        ans = (1LL * ans + 1LL * a[i] * tot) % mod;
    }
    cout<<ans<<'\n';
}

同类题：

结合数论分块

LightOJ 1098. A New Function

同类题：

蓝桥杯2022年第十三届省赛真题-因数平方和

结合单调栈

6077. 巫师的总力量和 - 力扣（LeetCode）

E - Second Sum (atcoder.jp)（单调栈求左右第二个比自己大的数）

同类题：

算法
单调栈

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议 ，转载请注明出处！