题解：P14514 [NFLSPC #8] 如何区分北京东路和北京东路

· · 题解

题意简述

给定长度为 n 的数列 a_i，每次随机选择 1 个元素平均分给其他 n-1 个元素，求 k 次操作后 a_i 的期望。

解题思路

数列 a_i 的总和始终为定值 s：

s=\sum_{i=1}^n a_i

设 j\in[0,k] 次操作后 a_i 的期望为 a_{i,j}。特别地，a_{i,0}=a_i。

如果 j 次操作选择了 p\in[1,n]，则有：

a_{i,j}= \begin{cases} 0 & p=i \\ a_{i,j-1}+\frac{a_{p,j-1}}{n-1} & p\ne i \end{cases}

因此 a_{i,j} 的期望为：

\begin{aligned} a_{i,j} &= \sum_{p=1,p\ne i}^n\frac{1}{n}\left(a_{i,j-1}+\frac{a_{p,j-1}}{n-1}\right) \\ &= \frac{n-1}{n}\cdot a_{i,j-1}+\frac{1}{n(n-1)}\sum_{p=1,p\ne i}^n a_{p,j-1} \\ &= \frac{n-1}{n}\cdot a_{i,j-1}+\frac{s-a_{i,j-1}}{n(n-1)} \\ &= \frac{n-1}{n}\cdot a_{i,j-1}+\frac{s}{n(n-1)}-\frac{1}{n(n-1)}\cdot a_{i,j-1} \\ &= \frac{(n-1)^2-1}{n(n-1)}\cdot a_{i,j-1}+\frac{s}{n(n-1)} \\ &= \frac{n-2}{n-1}\cdot a_{i,j-1}+\frac{s}{n(n-1)} \end{aligned}

显然 a_{i,j} 的值只与 a_{i,j-1} 有关，这是一个线性递推。设：

x=\frac{n-2}{n-1},y=\frac{s}{n(n-1)}

得到 a_{i,k} 的通项公式：

\begin{aligned} a_{i,k} &= x^ka_{i,0}+y\left(x^{k-1}+x^{k-2}+\dots+x+1\right) \\ &= x^k a_{i,0}+y\cdot\frac{x^k-1}{x-1} \end{aligned}

参考代码

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

using ll=long long;
const int mod=998244353;
const int N=1000005;
ll a[N];
ll Pow(ll x,ll y)
{
    x%=mod;
    ll res=1;
    while(y)
    {
        if(y&1)res=res*x%mod;
        x=x*x%mod;
        y>>=1;
    }
    return res;
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n,k;
    cin>>n>>k;
    ll sum=0;
    for(int i=1;i<=n;i++)
    {
        cin>>a[i];
        sum=(sum+a[i])%mod;
    }
    ll x=(n-2)*Pow(n-1,mod-2)%mod,y=sum*Pow(n,mod-2)%mod*Pow(n-1,mod-2)%mod;
    for(int i=1;i<=n;i++)
    {
        cout<<(Pow(x,k)*a[i]%mod+y*(Pow(x,k)-1)%mod*Pow(x+mod-1,mod-2)%mod)%mod<<' ';
    }
    return 0;
}