Official

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\) で割れば答えになります。

アルゴリズム

  1. 部員を技術力の昇順にソートします。

  2. 体力の値を重複を除いて昇順に並べた配列 values を作ります。

    • values の各要素に、\(1\) から始まる番号を割り当てます。
    • Fenwick Tree には、その番号に対応する体力を持つ候補の人数を記録します。
  3. 次の変数を初期化します。

    • p:まだ追加していない部員のうち、技術力が最大の位置
    • inserted:追加済みの候補の総人数
    • answer:条件を満たす相手の人数の累計
  4. ソート順に、各部員 \((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\) を引きます。

  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 によって生成されました。