ABC276 F問題（Double Chance）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）276 のF問題をC++で解いてみました。ABC276は、2022年11月5日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

F問題 Double Chance（Difficulty : 1562）

問題はリンク先をご覧ください。

ABC276 F問題 Double Chance

セグメント木を二つ使うことにより解くことができます。AtCoder Problems による Difficulty は、1562 でした。

解答案

$K$ を固定した場合、2枚のカードを引く方法は $K^2$ 通りあり、それぞれが同じ確率で起こります。求める確率は、以下となります。

$$(\sum _{i=1}^K\sum _{j=1}^{K}max(A_i,A_j))/K^2$$

この分子の式を展開すると、以下となります。

$$\sum _{i=1}^K\sum _{j=1}^{K}max(A_i,A_j)=\sum _{i=1}^{K-1}\sum _{j=1}^{K-1}max(A_i,A_j)+2\sum _{i=1}^{K-1}max(A_i,A_K)+A_K$$

最初の項は、$K=1$ から順に計算していくことで求めることができます。2番目の項（の半分）は以下となります。

$$\sum _{i=1}^{K-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 への移植版のプログラムの紹介を省略します。

最後に

コンテスト中には解けませんでしたが、セグメント木を使うと解けました。理解を深めるために記事にしました。

類題は、ABC351F問題（解説記事）です。

引き続き ABC の問題を紹介していきます。