E - ペアの選定 / Choosing Pairs Editorial by admin
gpt-6-astra-medium概要
技術力の合計が \(X\) 以上、体力の合計が \(Y\) 以上になる部員のペアを数えます。技術力でソートして候補を順に追加し、Fenwick Tree(BIT)で体力の条件を満たす人数を高速に求めます。
考察
全ペアの確認では間に合わない
すべての組 \((i,j)\) を調べると、時間計算量は \(O(N^2)\) です。\(N=10^5\) の場合、約 \(50\) 億組を確認する必要があり、実行時間に間に合いません。
そこで、部員を \(1\) 人固定し、条件を満たす相手の人数をまとめて数えることを考えます。
技術力が \(a\)、体力が \(b\) の部員に対して、相手に必要な条件は次の通りです。
- 技術力が \(X-a\) 以上
- 体力が \(Y-b\) 以上
技術力の条件を満たす候補は、追加するだけで管理できる
部員を技術力の昇順にソートし、その順番に注目する部員を処理します。
注目する部員の技術力 \(a\) が大きくなると、相手に必要な技術力の下限 \(X-a\) は小さくなります。つまり、技術力の条件を満たす候補は増えるだけで、いったん追加した候補を取り除く必要はありません。
そこで、ソート済みの配列の末尾からポインタを動かし、技術力が \(X-a\) 以上の部員を順にデータ構造へ追加します。
各部員は高々 \(1\) 回しか追加されないため、候補を追加する回数の合計は \(O(N)\) です。
体力の条件は Fenwick Tree で数える
候補のうち、体力が \(Y-b\) 以上の人数は、
\(候補の総人数 - 体力が Y-b 未満の候補人数\)
で求められます。
Fenwick Tree は、値ごとの人数を管理し、「先頭からある位置までの人数の合計」を \(O(\log N)\) で求められるデータ構造です。これを使えば、体力が一定値未満の候補人数を高速に数えられます。
ただし、体力は最大 \(10^9\) なので、その値をそのまま配列の添字にはできません。体力の値を小さい順に並べ、それぞれを連続した番号に置き換える座標圧縮を行います。
自分自身と二重計上に注意する
この方法では、候補に自分自身が含まれる場合があります。
自分自身が数えられる条件は、
\(2a \ge X\) かつ \(2b \ge Y\)
です。このときは、数えた人数から \(1\) を引きます。
また、異なる \(2\) 人の有効なペアは、どちらを注目する部員にするかによって \(2\) 回数えられます。そのため、最後に合計を \(2\) で割れば答えになります。
アルゴリズム
部員を技術力の昇順にソートします。
体力の値を重複を除いて昇順に並べた配列
valuesを作ります。valuesの各要素に、\(1\) から始まる番号を割り当てます。- Fenwick Tree には、その番号に対応する体力を持つ候補の人数を記録します。
次の変数を初期化します。
p:まだ追加していない部員のうち、技術力が最大の位置inserted:追加済みの候補の総人数answer:条件を満たす相手の人数の累計
ソート順に、各部員 \((a,b)\) について次を行います。
① 技術力の条件を満たす部員を追加する
members[p] の技術力が \(X-a\) 以上である間、その部員の体力に対応する Fenwick Tree の位置へ \(1\) を加えます。その後、inserted を増やし、p を左へ動かします。
この処理後、Fenwick Tree には、技術力が \(X-a\) 以上の部員がちょうど登録されています。
② 体力の条件を満たす人数を求める
bisect_left(values, Y-b) により、体力が \(Y-b\) 未満である値の種類数 k を求めます。
Fenwick Tree の先頭 k 個の合計を below とすると、両方の条件を満たす人数は inserted - below です。これを answer に加えます。
例えば、候補の体力が \([2,5,5,9]\) で、相手に必要な体力が \(5\) 以上なら、\(5\) 未満は \(1\) 人なので、条件を満たす人数は \(4-1=3\) 人です。
③ 自分自身を除外する
\(2a \ge X\) かつ \(2b \ge Y\) なら、自分自身も数えられているため answer から \(1\) を引きます。
- 各ペアは \(2\) 回ずつ数えられているので、
answer // 2を出力します。
計算量
- 時間計算量: \(O(N\log N)\)
- 空間計算量: \(O(N)\)
ソートと座標圧縮に \(O(N\log N)\) かかります。Fenwick Tree への追加は全体で \(N\) 回、人数の問い合わせも \(N\) 回で、それぞれ \(O(\log N)\) です。
実装のポイント
Fenwick Tree の添字は \(1\) 始まりにする
コードではranksによって、体力を \(1\) 以上の添字へ変換しています。「未満」と「以下」を区別する
求めたいのは体力が \(Y-b\) 以上の人数なので、引くのは \(Y-b\) 未満の人数です。bisect_leftを使うことで、ちょうど \(Y-b\) の部員を正しく残せます。二分探索の戻り値を、そのまま区間和の範囲に使える
k = bisect_left(values, Y-b)は、条件を満たさない値がvaluesの先頭に何種類あるかを表します。それらの Fenwick Tree 上の添字は \(1\) からkなので、先頭k個の合計を取得します。同じ能力の部員も別々に数える
座標圧縮で体力の値をまとめても、Fenwick Tree には部員ごとに \(1\) を加えるため、人数は正しく管理できます。答えには大きな整数型が必要
答えは最大で \(\frac{N(N-1)}{2}\) となり、\(32\) ビット整数には収まりません。Python では問題ありませんが、他の言語では \(64\) ビット整数型を使用します。ソースコード
import sys
from bisect import bisect_left
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
n, x, y = data[:3]
members = list(zip(data[3::2], data[4::2]))
del data
members.sort()
values = sorted({b for _, b in members})
m = len(values)
ranks = {b: i + 1 for i, b in enumerate(values)}
bit = [0] * (m + 1)
p = n - 1
inserted = 0
answer = 0
for a, b in members:
threshold = x - a
while p >= 0 and members[p][0] >= threshold:
k = ranks[members[p][1]]
while k <= m:
bit[k] += 1
k += k & -k
inserted += 1
p -= 1
k = bisect_left(values, y - b)
below = 0
while k:
below += bit[k]
k -= k & -k
answer += inserted - below
if a + a >= x and b + b >= y:
answer -= 1
print(answer // 2)
if __name__ == "__main__":
main()
この解説は gpt-6-astra-medium によって生成されました。