AtCoder が提供しているABC（AtCoder Beginner Contest）331 のC問題をC++とPythonで解いてみました。ABC331は、2023年12月2日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
C問題 Sum of Numbers Greater Than Me（Difficulty : 288）
問題はリンク先をご覧ください。
ABC331 C問題 Sum of Numbers Greater Than Me
ソートして累積和を求めておきます。AtCoder Problems による Difficulty は 288 でした。
解答案
数列の長さ $N$ の制約は、$1 \leqq N \leqq 2 \times 10^5$ です。それぞれの a[i] に対して、それより大きな要素を求めるために $N$ 回の比較が必要となり、合わせて計算量は、$N^2 = 4 \times 10^{10}$ となり間に合いません。
ソートして累積和を求めることにより計算量を下げます。
C++ プログラム例（ABC331C）
配列 a をソートした配列を b とします（15行目）。この配列 b の累積和 s を求めておきます（16ー20行目）。この累積和 s を使うと、[left, right) の半開区間に含まれる要素の全ての和は、s[right] – s[left] で求めることができます。
a[i] より大きな要素の場所を upper_bound で求めます。上記の式から、a[i]より大きな要素全ての和が求まります（23、24行目）。
計算量は、二分探索と合わせて、$O(N \log N)$ と下げることができました。以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
b[i] = a[i];
}
sort(b.begin(), b.end());
vector<ll> s(n + 1);
s[0] = 0;
for (int i = 0; i < n; ++i) {
s[i + 1] = s[i] + b[i];
}
for (int i = 0; i < n; ++i) {
auto it = upper_bound(b.begin(), b.end(), a[i]);
cout << s[n] - s[it - b.begin()] << " \n"[i == n - 1];
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC331C）
Python の bisect.bisect_right は、C++ の upper_bound に相当します（13行目）。結果はリスト result に格納して、まとめて出力しています（15行目）。
"""AtCoder Beginner Contest 331 C"""
import bisect
n = int(input())
a = list(map(int, input().split()))
b = sorted(a)
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = s[i] + b[i]
result = []
for i in range(n):
result.append(s[n] - s[bisect.bisect_right(b, a[i])])
print(*result)こちらも「AC」と判定されました。
最後に
事前に累積和を計算しておくことにより、計算量を下げることができました。
引き続き ABC の問題を紹介していきます。
