F - Double Chance | ABC276 | ABC | AtCoder | 競プロメモ (beta)

問題

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

考察

期待値

まず、期待値の分母は $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$ をかける。

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

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

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

\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}

ソート済みの制約を外す

さすがに「 $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;
}