F - Double Chance

cpatcodermathexpected valuediff cyan

最終更新日

問題

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

問題文

NN 枚のカードそれぞれに整数 AiA_i が書かれている。

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

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

制約

サンプル

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

考察

期待値

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

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

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

∑i=1KAi⋅(2i−1)\displaystyle \sum_{i=1}^K A_i \cdot (2 i - 1)

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

{k=1→E=A1k=2→E=A1+3A24k=3→E=A1+3A2+5A39\begin{cases} & k = 1 \rightarrow E = A_1 \\ & k = 2 \rightarrow E = \frac{A_1 + 3 A_2}{4} \\ & k = 3 \rightarrow E = \frac{A_1 + 3 A_2 + 5 A_3}{9} \end{cases}

ソート済みの制約を外す

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

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

例えば、 K=4→5,A=(1,2,4,5,3)K=4 \rightarrow 5, A = (1, 2, 4, 5, 3) のとき、分子だけ見た E=1+3⋅2+5⋅4+7⋅5→1+3⋅2+5⋅3+7⋅4+9⋅5E = 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⋅35 \cdot 3 と 2⋅(4+5)2 \cdot (4 + 5) を足せば OK。

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

データ構造

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

高速に求めたいのは、「AKA_K が何番目の数か？」「i+1,...,Ki+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;
}