ABC276 F - Double Chance を解くときに考えたこと

はじめに

自分の思考を殴り書いているような面があるので、文章の体裁はあまり保たれていませんし、解説もそこまでわかりやすくないです。
また、厳密性を欠くような部分もあると思いますが、ご容赦ください。

問題概要

整数列 $\lbrace A_i \rbrace _{(1 \leq i \leq N)},(1 \leq N \leq 2 \times 10^{5},1 \leq A_i \leq 2 \times 10^{5})$ が与えられます。
$1 \leq K \leq N$ について、以下のクエリに答えてください。

$A$ の部分列 $\lbrace A_i \rbrace _{1\leq i \leq K}$ から無作為に $1$ つの要素を選ぶことを二回繰り返します。(このとき、要素が重複してもよいです)
選ばれた要素のうち最大のものの期待値を有理数 $\mod 998244353$ で出力してください。

とっかかり

とりあえずある $K$ に対する期待値がいくらか考えてみます。
$\lbrace A_i \rbrace _ {(1 \leq i \leq K)}$ がソートされていると嬉しいので、例えば std::multiset などに入れ、ソートされた $\lbrace A_i \rbrace _ {(1 \leq i \leq K)}$ を考えたいです。
混同を防ぐためにソートされた $\lbrace A_i \rbrace _ {(1 \leq i \leq K)}$ を $\lbrace B _ {K _ i} \rbrace _ {(1 \leq i \leq K)}$ と書くことにします。
最初に $B _ {K _ i}$ が選ばれたとき、それが最後まで最大である場合の数を考えてみましょう。
これは二回目で $B _ {K _ i}$ 以下の値を引けばいいですから、 $B _ {K _ i}$ を超えない最大のインデックスを $j _ {K _ i}$ として $j _ {K _ i}$ です。
逆に、二回目で $B _ {K _ i}$ が選ばれたとき、それが最大である場合の数も、 $j _ {K _ i}$ です。
よって、 $B _ {K _ i}$ が最大である場合の数は $2j _ {K _ i}$ であるように思えますが、罠があります。
それは $B _ {K _ i}$ を二回選ぶような方法を重複してカウントしていることです。
よってその分を引いて場合の数は $2j _ {K _ i}-1$ であると分かりました。
したがって求めるべき期待値は $\displaystyle \frac{1}{K^{2}} \sum _ {i=1}^{K} (2j _ {K _ i} -1)B _ {K _ i}$ です。
これを便宜上、 $C _ {K}$ と書くことにします。

計算量改善

さて、これを愚直に計算すると、全体計算量 $O(N^{2})$ になってしまい間に合いません。(ユーザー解説にある通りコンパイル設定をいじれば通せはします)
そこで計算結果を使いまわすことを考えます。
$A _ {K+1}$ が、ソートした結果 $B _ {K+1 _ x}$ に入ったとして、期待値を計算しましょう。
$B _ {K+1 _ x}$ の期待値に対する寄与を考えると、これは $j _ {K+1 _ x}$ が求められれば良いです。
例えば std::multiset::upper_bound 等で $O(\log{K})$ 時間でできます。
次に、 $B _ {K+1 _ x}$ 以外の寄与を考えます。 $1 \leq i \lt x$ であるとき、 $j _ {K +1 _ i} = j _ {K _ i}$ であり、 $x \lt i \leq K+1$ のとき、 $j _ {K+1 _ i} = j _ {K _ {i-1}} +1$ です。
したがって $B _ {K+1 _ x}$ 以外の寄与は

$\displaystyle \dfrac{1}{(K+1)^2} \left(\sum _ {i = 1}^{x-1} (2j _ {K _ i}-1)B _ {K _ i} + \sum _ {i = x+1} ^ {K+1}(2j _ {K _ {i - 1}} + 1) B _ {K _ {i-1}}\right)$ $\displaystyle = \dfrac{1}{(K+1)^2} \left(\sum _ {i = 1}^{K} (2j _ {K _ i}-1)B _ {K _ i} + \sum _ {i = x} ^ {K} 2B _ {K _ i}\right)$ $\displaystyle = \dfrac{1}{(K+1)^2} \left(C _ K + \sum _ {i = x} ^ {K} 2B _ {K _ i}\right)$

と計算できます。
$C _ K$ は使いまわせますから、 $\displaystyle \sum _ {i = x} ^ {K} 2B _ {K _ i}$ を素早く計算することを考えます。区間和なのでセグ木を使いたいな～という気持ちになります。
そこで、重複している場合にインデックスを増やす座圧を行ったうえで、全要素 $0$ で初期化した長さ $N$ のセグ木を作りそこに要素を入れていけばよいです。
こうすると $\displaystyle \sum _ {i = x} ^ {K} 2B _ {K _ i}$ が $O(\log {K})$ で求められました。したがって一クエリあたりの計算量は $O(\log {K})$ であり、総計算量 $O(N \log{N})$ でこの問題に答えることができました。これは十分高速に動作し、 AC を得ることができます。