题解 | #小红的数组权值#

题目链接

小红的数组权值

题目描述

小红定义一个长度为 $m$ 的数组 $b$ 的权值为 $\sum_{k=1}^{m} k \cdot b_k$ 。

现在给定一个长度为 $n$ 的数组 $a$ ，她想知道所有子数组的权值和是多少？答案对 $10^9 + 7$ 取模。

解题思路

这是一个求解所有子数组特定属性之和的问题。直接枚举所有子数组并计算权值的 $O(N^2)$ 或 $O(N^3)$ 算法会超时，因此我们需要寻找一种更高效的计算方法。

贡献法

解决这类问题的常用技巧是贡献法。我们不去计算每个子数组的权值，而是反过来，计算原始数组中的每一个元素 $a_i$ 对最终的总和贡献了多少。最终答案就是所有元素贡献的总和。

推导单个元素的贡献

让我们来分析单个元素 $a_i$ （使用 0-based 索引, $i$ 从 $0$ 到 $n-1$ ）的总贡献。

$a_i$ 何时产生贡献？

当且仅当一个子数组包含 $a_i$ 时， $a_i$ 才会对该子数组的权值有贡献。一个子数组 a[l...r] 包含 $a_i$ 的充要条件是 $l \le i \le r$ 。
$a_i$ 在子数组中的贡献

在子数组 a[l...r] 中， $a_i$ 是第 $(i - l + 1)$ 个元素（1-based 索引）。因此，它对这个子数组权值的贡献是 $a_i \times (i - l + 1)$ 。
$a_i$ 的总贡献

$a_i$ 的总贡献是它在所有包含它的子数组中的贡献之和。

$\text{Contribution}(a_i) = \sum_{l=0}^{i} \sum_{r=i}^{n-1} a_i \times (i - l + 1)$
化简公式
- 将 $a_i$ 提出来：
  
  $\text{Contribution}(a_i) = a_i \times \left( \sum_{l=0}^{i} \sum_{r=i}^{n-1} (i - l + 1) \right)$
- 括号中的系数部分，内层求和项 $(i-l+1)$ 与 $r$ 无关。对于一个固定的 $l$ ，终点 $r$ 有 $n-i$ 种选择（从 $i$ 到 $n-1$ ）。因此，内层求和的结果是 $(n-i) \times (i-l+1)$ 。
- 公式变为：
  
  $\text{Contribution}(a_i) = a_i \times \left( \sum_{l=0}^{i} (n-i) \times (i-l+1) \right)$
- 再提出与 $l$ 无关的 $(n-i)$ ：
  
  $\text{Contribution}(a_i) = a_i \times (n-i) \times \left( \sum_{l=0}^{i} (i-l+1) \right)$
- 最后的求和部分 $\sum_{l=0}^{i} (i-l+1)$ 是 $(i+1) + i + \dots + 1$ 的和，即 $\frac{(i+1)(i+2)}{2}$ 。
- 最终，我们得到 $a_i$ 对总权值和的完整贡献：
  
  $\text{Contribution}(a_i) = a_i \times \frac{(n-i)(i+1)(i+2)}{2}$

最终算法

总权值和就是所有元素贡献的和：

$\text{Total Weight} = \sum_{i=0}^{n-1} \left( a_i \times \frac{(n-i)(i+1)(i+2)}{2} \right)$

这个公式可以在 $O(N)$ 的时间内完成计算。在实现时，所有计算都需要对 $10^9 + 7$ 取模，除以 2 的操作要用乘以 2 的模逆元（即 $500000004$ ）来代替。

代码

cpp
java
python

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

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;

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

    long long total_weight_sum = 0;
    long long MOD = 1e9 + 7;
    long long inv2 = 500000004;

    for (int i = 0; i < n; ++i) {
        long long val = a[i];
        long long n_ll = n;
        long long i_ll = i;

        long long term1 = (n_ll - i_ll);
        long long term2 = (i_ll + 1);
        long long term3 = (i_ll + 2);
        
        long long coeff = (term1 % MOD * (term2 % MOD)) % MOD;
        coeff = (coeff * (term3 % MOD)) % MOD;
        coeff = (coeff * inv2) % MOD;
        
        long long contribution = (val % MOD * coeff) % MOD;
        
        total_weight_sum = (total_weight_sum + contribution) % MOD;
    }

    cout << total_weight_sum << endl;

    return 0;
}

import java.util.Scanner;

public class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        long[] a = new long[n];
        for (int i = 0; i < n; i++) {
            a[i] = sc.nextLong();
        }

        long totalWeightSum = 0;
        long MOD = 1_000_000_007;
        long inv2 = 500000004;

        for (int i = 0; i < n; i++) {
            long val = a[i];
            long n_ll = n;
            long i_ll = i;

            long term1 = n_ll - i_ll;
            long term2 = i_ll + 1;
            long term3 = i_ll + 2;

            long coeff = (term1 % MOD * (term2 % MOD)) % MOD;
            coeff = (coeff * (term3 % MOD)) % MOD;
            coeff = (coeff * inv2) % MOD;

            long contribution = (val % MOD * coeff) % MOD;
            
            totalWeightSum = (totalWeightSum + contribution) % MOD;
        }

        System.out.println(totalWeightSum);
    }
}

import sys

def solve():
    try:
        n = int(sys.stdin.readline())
        a = list(map(int, sys.stdin.readline().split()))
    except (IOError, ValueError):
        return

    total_weight_sum = 0
    MOD = 10**9 + 7
    inv2 = 500000004

    for i in range(n):
        val = a[i]
        
        # Using 0-based index i
        term1 = n - i
        term2 = i + 1
        term3 = i + 2
        
        coeff = (term1 * term2) % MOD
        coeff = (coeff * term3) % MOD
        coeff = (coeff * inv2) % MOD
        
        contribution = (val * coeff) % MOD
        
        total_weight_sum = (total_weight_sum + contribution) % MOD
        
    print(total_weight_sum)

solve()

算法及复杂度

算法：贡献法、数学推导
时间复杂度： $O(N)$ 。
- 我们只需要遍历一次数组，每次循环中的计算都是常数时间的操作。
空间复杂度： $O(N)$ 。
- 需要 $O(N)$ 的空间来存储输入的数组。如果输入可以边读边处理，则空间复杂度可以降为 $O(1)$ 。