Official

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

gpt-5.5-high

概要

学生を頂点、\(\gcd(W_i, W_j) \geq K\) となるペアを辺とみなしたグラフを考え、その連結成分ごとの \(W_i\) の合計の最大値を求める問題です。

考察

素朴に全ての学生のペアについて \(\gcd(W_i, W_j)\) を計算すると、ペア数は \(O(N^2)\) です。
\(N \leq 2 \times 10^5\) なので、これは明らかに間に合いません。

ここで重要な観察は次です。

ある \(2\) つの値 \(x, y\) について、

\[ \gcd(x, y) \geq K \]

であることは、

\(K\) 以上のある整数 \(d\) が存在して、\(x\)\(y\) の両方を割り切る」

と言い換えられます。

つまり、ある \(d \geq K\) に対して、\(d\) の倍数である値同士は同じ連結成分に属します。

例えば \(K = 3\)、値が \(6, 10, 15\) の場合を考えます。

  • \(6\)\(15\) はともに \(3\) の倍数なので繋がる
  • \(10\)\(15\) はともに \(5\) の倍数なので繋がる
  • \(6\)\(10\) は直接は繋がらないが、\(15\) を経由して間接的に繋がる

したがって、これら \(3\) つは同じグループになります。

このように、「値」を頂点として、共通の約数 \(d \geq K\) を持つ値同士を Union-Find でまとめていけば、連結成分を求められます。

また、同じ値が複数回出てくる場合にも注意が必要です。

  • \(w \geq K\) の場合
    同じ値同士の \(\gcd\)\(w\) なので、同じ値を持つ学生たちは互いに繋がります。
    そのため、最初から合計 \(cnt[w] \times w\) の成分として扱えます。

  • \(w < K\) の場合
    どの値との \(\gcd\) も高々 \(w\) なので、\(K\) 以上にはなりません。
    つまり、同じ値が複数人いても互いに繋がらず、それぞれ独立した成分になります。
    この場合の成分の総合力は \(w\) です。

アルゴリズム

まず、各値 \(w\) の出現回数を cnt[w] として数えます。
最大値を \(M = \max W_i\) とします。

特殊ケース

\(K = 1\) の場合

任意の \(2\) つの正整数の最大公約数は必ず \(1\) 以上です。

したがって全ての学生が同じ連結成分になり、答えは

\[ \sum_i W_i \]

です。

\(K > M\) の場合

どの \(2\) つの値についても、最大公約数は高々 \(M\) です。
そのため \(\gcd(W_i, W_j) \geq K\) となるペアは存在しません。

全員が独立した成分になるので、答えは最大の専門スコア \(M\) です。

通常ケース

Union-Find を値 \(1, 2, \ldots, M\) に対して用意します。

各値 \(i\) について初期状態を設定します。

  • \(i < K\) かつ cnt[i] > 0 の場合
    その値を持つ学生は全員孤立しているので、答え候補は \(i\) です。

  • \(i \geq K\) かつ cnt[i] > 0 の場合
    同じ値を持つ学生同士は繋がるので、その成分の合計値を

$\( cnt[i] \times i \)$

として初期化します。

次に、\(d = K, K+1, \ldots, M\) について、\(d\) の倍数を調べます。

\(d\) の倍数で、かつ実際に入力に存在する値を全て同じ成分にまとめます。

具体的には、

d, 2d, 3d, ...

を順に見て、存在する最初の値を基準 base とし、以降に見つかった存在する倍数を Union します。

これにより、「共通の約数 \(d \geq K\) を持つ値同士」が同じ成分になります。

Union するときには、成分の合計値も一緒に管理します。
成分を併合したら、新しい合計値で答えを更新します。

この処理を全ての \(d \geq K\) について行えば、条件を満たす連結成分が全て正しく求められます。

計算量

\(M = \max W_i\) とします。

\(d\) についてその倍数を調べるので、全体のループ回数は

\[ \sum_{d=K}^{M} \frac{M}{d} \]

程度です。これは高々 \(O(M \log M)\) です。

Union-Find の操作はほぼ定数時間なので、全体として十分高速です。

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

実装のポイント

Union-Find では、各根に対してその連結成分の総合力を comp[root] として持ちます。

併合時には、

comp[new_root] = comp[root1] + comp[root2]

のように合計値を更新し、その値で答えも更新します。

また、値 \(i < K\) は他のどの値とも繋がれない点に注意が必要です。
特に、同じ値が複数回出てきても、\(i < K\) ならそれらは互いに繋がりません。
そのため cnt[i] * i を成分の合計にしてはいけません。

ソースコード

import sys

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    N, K = data[0], data[1]
    W = data[2:]

    M = max(W)

    if K == 1:
        print(sum(W))
        return

    if K > M:
        print(M)
        return

    cnt = [0] * (M + 1)
    for w in W:
        cnt[w] += 1

    parent = list(range(M + 1))
    comp = [0] * (M + 1)

    ans = 0

    for i in range(1, K):
        if cnt[i]:
            ans = i

    for i in range(K, M + 1):
        if cnt[i]:
            v = cnt[i] * i
            comp[i] = v
            if v > ans:
                ans = v

    p = parent
    s = comp
    c = cnt

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

    for d in range(K, M + 1):
        base = 0
        for m in range(d, M + 1, d):
            if c[m]:
                if base == 0:
                    base = find(m)
                else:
                    r = find(m)
                    if base != r:
                        if s[base] < s[r]:
                            base, r = r, base
                        p[r] = base
                        ns = s[base] + s[r]
                        s[base] = ns
                        s[r] = 0
                        if ns > ans:
                            ans = ns

    print(ans)

if __name__ == "__main__":
    main()

この解説は gpt-5.5-high によって生成されました。

posted:
last update: