公式

E - 研究グループの編成 / Formation of Research Groups 解説 by admin

gemini-3.5-flash-high

概要

この問題は、与えられた条件(最大公約数が \(K\) 以上)を満たす学生同士を同じグループ(連結成分)にまとめたとき、グループ内の学生の専門スコアの総和(総合力)の最大値を求める問題です。

グラフの連結成分を効率よく管理するために Union-Find(DSU) を使用し、約数・倍数の関係性を利用して高速にグループ分けを行います。


考察

1. ナイーブなアプローチとその限界

すべての学生のペア \((i, j)\) について \(\gcd(W_i, W_j) \ge K\) であるかを判定し、条件を満たすペアの間に辺を張る方法が最初に思い浮かびます。しかし、学生の数 \(N\) は最大で \(2 \times 10^5\) であるため、ペアの数は \(O(N^2) \approx 4 \times 10^{10}\) となり、実行時間制限に間に合いません(TLEとなります)。

2. スコアの範囲と「協力可能」の条件の言い換え

各学生のスコア \(W_i\) は最大で \(M = 10^6\) です。この数値の小ささに着目します。

  • \(W_i < K\) の学生 いかなる学生 \(j\) に対しても、\(\gcd(W_i, W_j) \le W_i < K\) となるため、誰とも協力できません。したがって、これらの学生は必ず「自分1人だけのグループ」になります。
  • \(W_i \ge K\) の学生 同じスコアを持つ学生同士は、最大公約数が自分自身のスコア(\(\ge K\))になるため、必ず同じグループに属します。

異なるスコア \(a, b \ge K\) を持つ学生同士が協力できるのは、\(\gcd(a, b) = g \ge K\) のときです。これは、\(a\)\(b\) が、ある \(g \ge K\) の倍数である」と言い換えることができます。 もし \(a, b\) がともに \(g\) の倍数であれば、\(\gcd(a, b)\) は必ず \(g\) の倍数(すなわち \(g\) 以上)になるため、彼らは直接協力可能です。

3. 調和級数を利用した高速なマージ

\(g \ge K\) について、\(g\) の倍数であるスコア(\(g, 2g, 3g, \ldots\))を持つ学生たちをすべて同じグループにマージすることを考えます。

\(g\)\(K\) から \(M\) まで全探索し、各 \(g\) について \(g\) の倍数を走査します。 このとき、走査する数の個数は以下のようになります。 $\( \frac{M}{K} + \frac{M}{K+1} + \dots + \frac{M}{M} \le M \left( 1 + \frac{1}{2} + \dots + \frac{1}{M} \right) \approx M \log M \)\( これは「調和級数」の性質より、全体で約 \)M \log M$ 回の計算量となり、十分に高速です。


アルゴリズム

  1. \(W_i < K\) の処理 誰ともグループを組めないため、これらのスコアの最大値をあらかじめ記録しておきます(初期の最大総合力候補)。
  2. スコアごとの総和の計算 \(W_i \ge K\) の学生について、同じスコアを持つ学生のスコアの総和 sum_W[w] を計算します。
  3. Union-Find の初期化 \(0\) から \(M\) までの要素を持つ Union-Find を用意します。各要素 \(w\) の初期の重み(グループの総合力)は sum_W[w] とします。
  4. 倍数のマージ \(g\)\(K\) から \(M\) までループします。
    • \(g\) の倍数 \(x = g, 2g, 3g, \ldots \le M\) を順番に見ます。
    • 実際に存在する(sum_W[x] > 0 である)スコア \(x\) を探します。
    • 最初に見つかったスコアを first とし、それ以降に見つかった存在する倍数 xfirst と同じグループにマージ(Union)していきます。マージする際、Union-Find のルートの重みを足し合わせます。
  5. 答えの出力 Union-Find の各グループの総合力の最大値と、ステップ1で記録した単独グループの最大値のうち、大きい方を出力します。

計算量

  • 時間計算量: \(O(N + M \log M)\)

    • \(W_i < K\) の処理、および sum_W の構築に \(O(N)\) かかります。
    • 倍数のマージにおけるループ回数は調和級数より \(O(M \log M)\) 回です。Union-Find の各操作(Find, Union)はほぼ定数時間 \(O(\alpha(M))\) で行えるため、マージ部分は \(O(M \log M)\) となります。
    • 全体として実行時間制限に余裕で間に合います。
  • 空間計算量: \(O(N + M)\)

    • スコアの最大値 \(M \le 10^6\) に対し、Union-Find の親配列や重み配列のサイズが \(O(M)\) となります。

実装のポイント

  • Union-Find の高速化(Python向け) Python では、再帰関数を用いた Union-Find の find 操作は呼び出しオーバーヘッドにより遅くなることがあります。 提示されたコードでは、find 操作を while ループを用いてインライン展開(関数呼び出しをせず直接記述)し、さらに経路圧縮もループ内で行うことで、大幅な高速化を図っています。

  • 無駄なマージのスキップ sum_W[x] > 0 である(=そのスコアを持つ学生が実際に存在する)場合のみマージ対象とすることで、不要な Union-Find の操作を避けています。

    ソースコード

import sys


def solve():
    # 高速な入力読み込み
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    N = int(input_data[0])
    K = int(input_data[1])
    W = [int(x) for x in input_data[2:]]

    M = max(W)

    # W_i < K の学生は誰とも協力できないため、単独のグループとなる
    # その中での最大スコアを初期値とする
    ans = 0
    for w in W:
        if w < K:
            if w > ans:
                ans = w

    # W_i >= K の学生について、同じスコアを持つ学生のスコアの総和を計算
    sum_W = [0] * (M + 1)
    for w in W:
        if w >= K:
            sum_W[w] += w

    # Union-Find(DSU)の初期化
    parent = list(range(M + 1))
    weight = list(sum_W)

    # 各 g >= K について、g の倍数同士をマージする
    for g in range(K, M + 1):
        first = -1
        root_first = -1
        for x in range(g, M + 1, g):
            if sum_W[x] > 0:
                if first == -1:
                    first = x
                    # find(x) のインライン展開による高速化
                    curr = x
                    while parent[curr] != curr:
                        curr = parent[curr]
                    node = x
                    while node != curr:
                        nxt = parent[node]
                        parent[node] = curr
                        node = nxt
                    root_first = curr
                else:
                    # find(x) のインライン展開による高速化
                    curr = x
                    while parent[curr] != curr:
                        curr = parent[curr]
                    node = x
                    while node != curr:
                        nxt = parent[node]
                        parent[node] = curr
                        node = nxt
                    root_x = curr

                    if root_first != root_x:
                        parent[root_x] = root_first
                        weight[root_first] += weight[root_x]
                        weight[root_x] = 0

    # 連結成分の総合力の最大値を求める
    max_weight = max(weight)
    if max_weight > ans:
        ans = max_weight

    print(ans)


if __name__ == "__main__":
    solve()

この解説は gemini-3.5-flash-high によって生成されました。

投稿日時:
最終更新: