AtCoder が提供しているABC（AtCoder Beginner Contest）276 のF問題をC++で解いてみました。ABC276は、2022年11月5日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
F問題 Double Chance（Difficulty : 1562）
問題はリンク先をご覧ください。
セグメント木を二つ使うことにより解くことができます。AtCoder Problems による Difficulty は、1562 でした。
解答案
$K$ を固定した場合、2枚のカードを引く方法は $K^2$ 通りあり、それぞれが同じ確率で起こります。求める確率は、以下となります。
$$\left(\sum^K_{i=1} \sum^K_{j=1} \max (A_i, A_j) \right) / K^2$$
この分子の式を展開すると、以下となります。
$$\sum^K_{i=1} \sum^K_{j=1} \max (A_i, A_j) = \sum^{K-1}_{i=1} \sum^{K-1}_{j=1} \max (A_i, A_j) + 2\sum^{K-1}_{i=1} \max(A_i, A_K) + A_K$$
最初の項は、$K=1$ から順に計算していくことで求めることができます。2番目の項（の半分）は以下となります。
$$\sum^{K-1}_{i=1} \max(A_i, A_K) = \sum_{1 \leqq i \leqq K-1, A_i \leqq A_K}A_K + \sum_{1 \leqq i \leqq K-1, A_i > A_K} A_i$$
右辺の最初の項は、$A_k$ より小さい $A_i$ の個数に $A_k$ を掛けた値になります。次の項は、$A_k$ より大きい $A_i$ の総和となります。どちらもセグメント木で解くことができます。
C++ プログラム例（ABC276F）
上で考察したようにセグメント木を2つ用意します。ACL を使いました。
- セグメント木 seg1 で、$A_K$ より大きい、$A_i$ の総和を求めます。
- セグメント木 seg2 で、$A_K$ より小さい、$A_i$ の個数を求めます。
- 998244353 との剰余を求めるためにセグメント木と同じく ACL を使いました。
ACL のおかげで全体がすっきりと書けています。以下が、C++プログラムとなります。
#include <bits/stdc++.h>
#include <atcoder/modint>
#include <atcoder/segtree>
using namespace std;
using namespace atcoder;
typedef long long int ll;
typedef modint998244353 mint;
ll op(ll a, ll b)
{
return a + b;
}
ll e()
{
return 0LL;
}
int main()
{
int n;
cin >> n;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
segtree<ll, op, e> seg1(200001);
segtree<ll, op, e> seg2(200001);
mint prev = (mint)0;
for (int i = 0; i < n; ++i) {
mint t = 0;
t += seg1.prod(a[i], 200001);
seg1.set(a[i], seg1.get(a[i]) + a[i]);
t += seg2.prod(0, a[i]) * a[i];
seg2.set(a[i], seg2.get(a[i]) + 1);
prev += 2 * t + (mint)a[i];
mint result = prev / (i + 1) / (i + 1);
cout << result.val() << endl;
}
return 0;
}このプログラムは、AC（Accepted＝正しいプログラム）と判定されました。
Python のセグメント木について
Python は、セグメント木を標準ライブラリとしては用意していません。このブログの Python プログラムは、なるべく言語標準のライブラリで実装したいと考えています。このような事情で、この記事では、Python への移植版のプログラムの紹介を省略します。
最後に
コンテスト中には解けませんでしたが、セグメント木を使うと解けました。理解を深めるために記事にしました。
引き続き ABC の問題を紹介していきます。
