Official

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

claude4.8opus-high

概要

\(N\) 人の学生を「\(\gcd(W_i, W_j) \geq K\) なら協力可能」という関係で結んだときの連結成分に分け、各成分の専門スコアの合計のうち最大値を求める問題です。

考察

まず、この問題は Union-Find(素集合データ構造) を使った連結成分の管理に帰着できます。協力可能なペアを辺と見なし、辺で繋がる頂点をまとめていけば、最終的な連結成分が各グループに対応します。

しかし、素朴に「すべてのペア \((i, j)\) について \(\gcd(W_i, W_j) \geq K\) を判定して辺を張る」方法では、ペアの数が \(O(N^2)\) 個あり、\(N \leq 2 \times 10^5\) では計算が間に合いません(TLE)。

ここで重要な観察をします。

観察1: 共通の約数に注目する

\(\gcd(W_i, W_j) \geq K\) ということは、「\(K\) 以上のある整数 \(d\) が存在して、\(d\)\(W_i\)\(W_j\) の両方を割り切る」ことと同値です。

つまり、\(K\) 以上の各 \(d\) について、「\(d\) の倍数を専門スコアに持つ学生たち」は全員同じグループに繋げてよい、ということになります。なぜなら、それらの学生は互いに \(d\) という共通の約数を持つため、\(\gcd \geq d \geq K\) となり協力可能だからです。

逆に、\(\gcd(W_i, W_j) \geq K\) となるなら、その \(\gcd\) の値(\(\geq K\))を \(d\) とすれば、\(d\) の倍数のグループで両者は繋がります。したがって、この方法ですべての協力可能関係を漏れなく表現できます。

観察2: 同じ \(W\) を持つ学生のまとめ方

同じ専門スコア \(W\) を持つ学生が複数いるとき、\(W \geq K\) ならば彼らは互いに協力可能(\(\gcd(W, W) = W \geq K\))なので、まとめて1つにできます。各スコア値について「代表となる学生1人」を覚えておけば十分です。

アルゴリズム

エラトステネスの篩(ふるい)の要領で、各約数 \(d\) について倍数を走査していきます。

  1. 代表の記録: 各スコア値 \(w\) について、その値を持つ学生の代表 rep[w] を1人記録します。同じ値の学生が複数いて、かつ \(w \geq K\) ならば、それらを Union でまとめます。

  2. 約数ごとの連結: \(d = K, K+1, \ldots, M\)\(M\) は最大スコア)の各 \(d\) について、\(d, 2d, 3d, \ldots\) と倍数を辿ります。スコアが「\(d\) の倍数」である学生の代表を見つけたら、それらをすべて一つの集合に Union します。

    • 篩のように \(d\) の倍数を \(d, 2d, 3d, \dots\) と見ていくことで、\(d\) の倍数のスコアを持つ学生を効率よく集められます。
  3. 合計の集計: 最後に、各学生を Union-Find の根ごとに集約し、各連結成分の専門スコアの合計を計算します。その最大値が答えです。

具体例として \(K=3\)、スコアが \(\{6, 9, 10\}\) の場合を考えます。 - \(d=3\) の倍数は \(6, 9\) → スコア \(6\)\(9\) の学生が連結(\(\gcd(6,9)=3 \geq 3\))。 - \(d\) を進めても \(10\)\(3\) 以上の共通約数で他と繋がりません。

結果、\(\{6, 9\}\) のグループ(合計 \(15\))と \(\{10\}\) のグループ(合計 \(10\))に分かれ、答えは \(15\) となります。

計算量

\(M\) を専門スコアの最大値(\(\leq 10^6\))とします。

  • 時間計算量: 篩部分は \(d=K\) から \(M\) まで各 \(d\) について \(M/d\) 回の走査を行うので、調和級数の和より \(O\!\left(\sum_{d=K}^{M} \frac{M}{d}\right) = O(M \log M)\) です。Union-Find の操作はほぼ定数時間(逆アッカーマン関数)なので、全体で \(O(M \log M + N)\) となります。
  • 空間計算量: 代表配列 rep\(O(M)\)、Union-Find に \(O(N)\) で、合計 \(O(M + N)\) です。

実装のポイント

  • \(\gcd\) を直接計算しない: 約数 \(d\) を主役にして篩で繋ぐことで、\(\gcd\) の計算を一切行わずに済ませているのが本質です。これにより \(O(N^2)\) のペア判定を回避しています。

  • 代表でまとめる重要性: 同じスコアを持つ学生をあらかじめ代表1人に集約しておくことで、篩の走査で扱う対象を「スコア値の種類」単位にでき、無駄な Union を防げます。

  • \(d\) の開始位置: 共通約数 \(d\)\(K\) 以上でなければならないため、篩のループは d = K から始めます。これにより条件 \(\gcd \geq K\) を自然に満たします。

  • Union-Find の経路圧縮: find 関数で経路圧縮を行うことで、操作を高速化しています。最終的な合計集計時も find(i) で根を求めて加算します。

    ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0]); k = int(data[1])
    W = [int(x) for x in data[2:2+n]]
    if n == 0:
        print(0)
        return
    M = max(W)
    parent = list(range(n))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a, b):
        ra = find(a); rb = find(b)
        if ra != rb:
            parent[ra] = rb

    rep = [-1] * (M + 1)
    for i in range(n):
        w = W[i]
        if rep[w] == -1:
            rep[w] = i
        else:
            if w >= k:
                union(i, rep[w])

    # sieve: for each d >= k, union all students whose W is multiple of d
    for d in range(k, M + 1):
        first = -1
        m = d
        while m <= M:
            r = rep[m]
            if r != -1:
                if first == -1:
                    first = r
                else:
                    union(r, first)
            m += d

    total = [0] * n
    ans = 0
    for i in range(n):
        r = find(i)
        total[r] += W[i]
        if total[r] > ans:
            ans = total[r]
    print(ans)

main()

この解説は claude4.8opus-high によって生成されました。

posted:
last update: