F - Double Chance

問題

https://atcoder.jp/contests/abc276/tasks/abc276_f (新しいタブで開く)

問題文

$N$ 枚のカードそれぞれに整数 $A_i$ が書かれている。

$K = 1, ..., N$ について、以下の問題に答えてね。

カード $1, ..., K$ のうちから、カードを 1 枚ずつ引き出して戻す。
書かれた整数を、取った順に $x, y$ とする。
$\max(x,y)$ の期待値 $\mathrm{mod} 998244353$ は？

制約

• $1 \le N \le 2 \times 10^5$
• $1 \le A_i \le 2 \times 10^5$

サンプル

I/O 1

3
5 7 5

5
499122183
443664163

I/O 2

7
22 75 26 45 72 81 47

22
249561150
110916092
873463862
279508479
360477194
529680742

考察

期待値

まず、期待値の分母は $K^2$ でいい。これは選び方の全通りなので。

分子は？
仮に $A$ が昇順ソートされているとする。
$x = A_i, y = A_j$ として $x = \max(x,y)$ が成り立つとき、昇順なので $i \ge j$ が成り立つはず。逆も同様。
よって $A_i$ が max になる場合の数は $2 i - 1$
($x, y$で対称なので 2 をかける、$i = j$ を 2 回計上しているので -1)
期待値なのでその係数に $A_i$ をかける。

よって、分子部分を考えると、

具体例を挙げると、こんな感じ。

ソート済みの制約を外す

さすがに「 $k$ それぞれについてソートして計算」をすると $O(N^2 \log{N})$ で間に合わない。さらに考察が必要。

$K$ が 1 増えたときの動きを見てみる。とりあえず分子だけに注目する。
追加された数字が $i$ 番目に小さい数のとき、 $i+1, ..., K$ 番目の数は 2 倍を加えられる + $i \cdot A_K$ が加えられるっぽい。

例えば、 $K=4 \rightarrow 5, A = (1, 2, 4, 5, 3)$ のとき、分子だけ見た $E = 1 + 3 \cdot 2 + 5 \cdot 4 + 7 \cdot 5 \rightarrow 1 + 3 \cdot 2 + \textcolor{royalblue}{5 \cdot 3} + \textcolor{royalblue}{7} \cdot 4 + \textcolor{royalblue}{9} \cdot 5$ になってほしい。
こうしたとき、 $5 \cdot 3$ と $2 \cdot (4 + 5)$ を足せば OK。

これでとりあえず毎回ソートしなくても済みそう。

データ構造

データ構造どうする問題。

高速に求めたいのは、「$A_K$ が何番目の数か？」「$i+1, ..., K$ 番目の数の総和」の 2 つ。
フェニック木を使えばいい。 (← ここまでたどり着くのに本番 40 分かかった)

コード

オーバーフローには気を付けましょう (戒め) (新しいタブで開く)

https://atcoder.jp/contests/abc276/submissions/36259983 (新しいタブで開く)

int main() {
  int n;
  cin >> n;
  vector<ll> a(n);
  rep(i, n) cin >> a[i];

  map<ll, queue<ll>> mp, mp2;
  {  // 座圧、同じ数はまとめない
    ll cnt = 0;
    vector<ll> b = a;
    sort(b.begin(), b.end());
    rep(i, n) mp[b[i]].push(cnt++);
  }
  fenwick_tree<mint> fw(n);  // A_i の部分和求める用
  fenwick_tree<ll> fw2(n);  // 何番目の数か求める用
  mint now = 0;
  rep(i, n) {
    // これまでに出てきた数で、それ以下のものを数える
    ll idx = mp[a[i]].front();
    ll pushed = fw2.sum(0, idx);
    fw2.add(idx, 1);
    mp[a[i]].pop();
    // それより大きいものを2倍する
    now += fw.sum(idx, n) * 2;
    fw.add(idx, a[i]);
    // 適切な係数をかけて加える
    now += a[i] * (1 + 2 * pushed);
    cout << (now / mint(i + 1).pow(2)).val() << endl;
  }
  return 0;
}