E - Revenge of "The Salary of AtCoder Inc."

Problem Statement

Aoki, an employee at AtCoder Inc., has his salary for this month determined by an integer $N$ and a sequence $A$ of length $N$ as follows.
First, he is given an $N$-sided die (dice) that shows the integers from $1$ to $N$ with equal probability, and a variable $x=0$.

Then, the following steps are repeated until terminated.

• Roll the die once and let $y$ be the result.
  • If $x<y$, pay him $A_y$ yen and let $x=y$.
  • Otherwise, terminate the process.

Aoki's salary for this month is the total amount paid through this process.
Find the expected value of Aoki's salary this month, modulo $998244353$.

Constraints

• All inputs are integers.
• $1 \le N \le 3 \times 10^5$
• $0 \le A_i < 998244353$

Input

The input is given from Standard Input in the following format:

$N$
$A_1$ $A_2$ $\dots$ $A_N$

Output

Print the answer.

Sample Input 1

3
3 2 6

Sample Output 1

776412280

Here is an example of how the process goes.

• Initially, $x=0$.
• Roll the die once, and it shows $1$. Since $0<1$, pay him $A_1 = 3$ yen and let $x=1$.
• Roll the die once, and it shows $3$. Since $1<3$, pay him $A_3 = 6$ yen and let $x=3$.
• Roll the die once, and it shows $1$. Since $3 \ge 1$, terminate the process.

In this case, his salary for this month is $9$ yen.

It can be calculated that the expected value of his salary this month is $\frac{49}{9}$ yen, whose representation modulo $998244353$ is $776412280$.

Sample Input 2

1
998244352

Sample Output 2

998244352

Sample Input 3

9
3 14 159 2653 58979 323846 2643383 27950288 419716939

Sample Output 3

545252774

解题思路

  现在一遇到概率期望题就不会做......

  如果直接按照期望的公式来求，则需要把所有合法方案的权重以及对应的概率求出来，很明显这是很困难的。考虑贡献法，由于期望具有线性性，因此可以考虑每个 $a_i$ 对期望的贡献是多少，那么答案就是 $\sum\limits_{i=1}^{n}{a_i \cdot p_i}$，其中 $p_i$ 为获得能获得 $a_i$ 的概率。

  那么在哪些情况下能获得 $a_i$ 呢？分情况讨论：

• 如果第 $1$ 个数就是 $a_i$，则对应的概率为 $\frac{1}{n}$。
• 如果第 $2$ 个数是 $a_i$，那么前面的 $1$ 个数只能从 $[1, i-1]$ 中选，则对应的概率为 $\frac{C_{i-1}^{1}}{n} \cdot \frac{1}{n}$。
• 如果第 $3$ 个数是 $a_i$，那么前面的 $2$ 个数只能从 $[1, i-1]$ 中选，并且要升序排序，则对应的概率为 $\frac{C_{i-1}^{2}}{n^2} \cdot \frac{1}{n}$。
• $\ldots$
• 如果第 $k$ 个数是 $a_i$，那么前面的 $k-1$ 个数只能从 $[1, i-1]$ 中选，并且要升序排序，则对应的概率为 $\frac{C_{i-1}^{k-1}}{n^{k-1}} \cdot \frac{1}{n}$。
• $\ldots$
• 如果第 $i$ 个数是 $a_i$，那么前面的 $i-1$ 个数只能从 $[1, i-1]$ 中选，并且要升序排序，则对应的概率为 $\frac{C_{i-1}^{i-1}}{n^{i-1}} \cdot \frac{1}{n}$。

  因此 $p_i = \frac{1}{n} \cdot \sum\limits_{j=1}^{i}{\frac{C_{i-1}^{j-1}}{n^{j-1}}} = \frac{1}{n} \cdot \sum\limits_{j=0}^{i-1}{\frac{C_{i-1}^{j}}{n^{j}}}$。根据二项式定理有 $\frac{1}{n} \cdot \sum\limits_{j=0}^{i-1}{\frac{C_{i-1}^{j}}{n^{j}}} = \frac{1}{n} \cdot \left( 1 + \frac{1}{n} \right)^{i-1}$，因此所求的期望值就是 $\frac{1}{n} \cdot \sum\limits_{i=1}^{n}{a_i \cdot \left( 1 + \frac{1}{n} \right)^{i-1}}$。

  AC 代码如下，时间复杂度为 $O(n)$：

#include <bits/stdc++.h>
using namespace std;

const int N = 3e5 + 10, mod = 998244353;

int a[N];

int qmi(int a, int k) {
    int ret = 1;
    while (k) {
        if (k & 1) ret = 1ll * ret * a % mod;
        a = 1ll * a * a % mod;
        k >>= 1;
    }
    return ret;
}

int main() {
    int n;
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%d", a + i);
    }
    int ret = 0, p = qmi(n, mod - 2);
    for (int i = 1, t = 1; i <= n; i++) {
        ret = (ret + 1ll * a[i] * t) % mod;
        t = t * (1ll + p) % mod;
    }
    printf("%d", 1ll * ret * p % mod);
    
    return 0;
}

参考资料

  Editorial - Panasonic Programming Contest 2023（AtCoder Beginner Contest 326）：https://atcoder.jp/contests/abc326/editorial/7547

  AtCoder Beginner Contest 326 A 至 G 題讲解 by dreamoon：https://www.bilibili.com/video/BV1Ez4y1N7p1/

  AtCoder Beginner Contest 326 题解：https://www.cnblogs.com/binary1110011/p/17794737.html

  题解 ABC326E【Revenge of "The Salary of AtCoder Inc."：https://www.cnblogs.com/ruierqwq/p/abc326e.html