AtCoder が提供しているABC（AtCoder Beginner Contest）351 のF問題をC++で解いてみました。ABC351は、2024年4月27日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
F問題 Double Sum（Difficulty : 1253）
問題はリンク先をご覧ください。
セグメント木を二つ使うことにより解くことができます。AtCoder Problems による Difficulty は 1253 でした。
解答案
$i < j$ を満たして、$A_i < A_j$ となる場合だけが演算に影響を与えます。$A_i \geqq A_j$ の場合は、0を加えることになるため演算に影響を与えません。
$A_i$ を固定します。このときに、$i < j$ かつ $A_i < A_j$ となる $A_j$ の集合を特定することができて、$\sum_{i < j} (A_j \;-\; A_i)$ の値は以下の値から求めることができます。
- $A_i$ より大きい $A_j$ の個数
- $A_i$ より大きい $A_j$ の合計
この二つは、セグメント木で求めることができます。以下でプログラムを使って解説します。
C++ プログラム例（ABC351F）
セグメント木を二つ用意します。一方は、$A_i$ より値が大きい $A_j$ の合計を計算するために使います（st1）。他方は、$A_i$ より値が大きい $A_j$ の個数を計算するために使います（st2）。
このセグメント木は、値が大きい要素から処理をします。セグメント木に格納する位置は、配列と同じ場所とします。
値が大きい要素から処理をするために、値とインデックスの組をソートしておきます（23ー29行目）。
二つのセグメント木から [a[i], n) の値を取り出し、st1 からの値に合計として加えて、st2 からの値は、a[i] と掛けて引きます（35、36行目）。値を取り出した後に、セグメント木を更新します（37、38行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
#include <atcoder/segtree>
using namespace std;
using namespace atcoder;
typedef long long int ll;
ll op(ll a, ll b)
{
return a + b;
}
ll e(void)
{
return 0LL;
}
int main()
{
int n;
cin >> n;
vector<pair<ll, int>> a;
for (int i = 0; i < n; ++i) {
int t;
cin >> t;
a.push_back(make_pair(t, i));
}
sort(a.begin(), a.end(), greater<pair<ll, int>>());
ll result = 0;
segtree<ll, op, e> st1(n);
segtree<ll, op, e> st2(n);
for (int i = 0; i < n; ++i) {
result += st1.prod(a[i].second, n);
result -= st2.prod(a[i].second, n) * a[i].first;
st1.set(a[i].second, a[i].first);
st2.set(a[i].second, 1);
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python のセグメント木について
Python は、セグメント木を標準ライブラリとしては用意していません。このブログの Python プログラムは、なるべく言語標準のライブラリで実装したいと考えています。このような事情で、この記事では、Python への移植版のプログラムの紹介を省略します。
最後に
コンテスト中は、反転数と似た問題と考えて、遅延セグメント木と格闘していました。セグメント木を二つ使うことを公式解説でチラ見して、解くことができました。セグメント木は、水Diffまたは青Diffでよく出題される印象があります。
反転数（転倒数）を求める類題は、ABC190F問題（解説記事）です。
引き続き ABC の問題を紹介していきます。
