公式

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

gpt-5.5-high

概要

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

\(W_i\) の最大値が \(10^6\) と小さいことを利用し、「同じ約数 \(d \geq K\) を持つ値同士」を Union-Find でまとめます。

考察

素朴にすべての学生ペアについて \(\gcd(W_i, W_j)\) を計算すると、ペア数は最大で

\(O(N^2)\)

です。\(N \leq 2 \times 10^5\) なので、これは到底間に合いません。

そこで、学生そのものではなく「専門スコアの値」に注目します。

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

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

であることは、言い換えると

\(x\)\(y\) が共通して持つ約数 \(d\) が存在し、\(d \geq K\)

であるということです。

つまり、ある \(d \geq K\) に対して、\(d\) の倍数である値同士はすべて直接協力可能です。

例えば \(K = 3\) のとき、値 \(6, 15\) はどちらも \(3\) の倍数なので、

\(\gcd(6, 15) = 3 \geq K\)

となり、同じ成分に属します。

よって、各 \(d = K, K+1, \ldots\) について、入力に存在する \(d\) の倍数たちを Union-Find でまとめればよいです。

ただし、同じ値が複数回出現する場合に注意が必要です。

  • \(w \geq K\) の場合
    同じ値同士は \(\gcd(w, w) = w \geq K\) なので、同じグループになります。
    そのため、合計には \(w \times \text{freq}[w]\) を加えます。

  • \(w < K\) の場合
    任意の値 \(x\) に対して \(\gcd(w, x) \leq w < K\) なので、誰とも協力可能になりません。
    同じ値 \(w\) を持つ学生同士でも \(\gcd(w, w) = w < K\) なので別々のグループです。
    したがって、それぞれ単独のグループとなり、最大値候補は \(w\) です。

アルゴリズム

まず、各専門スコア \(w\) の出現回数を freq[w] に記録します。

また、専門スコアの最大値を \(A\) とします。

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

次に、各 \(d = K, K+1, \ldots, A\) について、\(d\) の倍数を列挙します。

for (int d = K; d <= A; d++) {
    int first = -1;
    for (int m = d; m <= A; m += d) {
        if (freq[m] == 0) continue;
        if (first == -1) first = m;
        else dsu.unite(first, m);
    }
}

ここで、\(d\) の倍数のうち入力に存在するものをすべて同じ集合にまとめます。

すべての倍数同士を直接 Union すると数が多くなりますが、代表となる first に対して順に Union すれば、同じ集合にまとめることができます。

例えば、\(d = 5\) の倍数として \(10, 15, 20\) が存在するなら、

  • \(10\)\(15\) を Union
  • \(10\)\(20\) を Union

すれば、\(10, 15, 20\) はすべて同じ集合になります。

その後、各値 \(w\) について合計を計算します。

  • \(w < K\) の場合は単独成分なので、答え候補は \(w\)
  • \(w \geq K\) の場合は Union-Find の根を調べ、その成分の合計に \(w \times \text{freq}[w]\) を足す

最後に、すべての成分合計の最大値を出力します。

計算量

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

\(d\) について倍数を列挙するので、列挙回数は

\(\sum_{d=K}^{A} \frac{A}{d} \leq A \log A\)

程度です。

Union-Find の操作はほぼ定数時間なので、

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

です。

制約では \(A \leq 10^6\) なので十分高速です。

実装のポイント

freq[w] に出現回数を持つことで、同じ値の学生をまとめて扱います。

ただし、\(w < K\) の値は同じ値同士でも協力可能ではないため、freq[w] 人をまとめてはいけません。
コードでは以下のように、\(w < K\) の場合は単独グループとして ans = max(ans, w) だけを行っています。

if (w < K) {
    ans = max(ans, (long long)w);
} else {
    int r = dsu.find(w);
    sum[r] += 1LL * w * freq[w];
}

また、合計値は最大で \(2 \times 10^5 \times 10^6 = 2 \times 10^{11}\) になるため、long long を使う必要があります。

ソースコード

#include <bits/stdc++.h>
using namespace std;

struct DSU {
    vector<int> parent, sz;

    DSU(int n) : parent(n), sz(n, 1) {
        iota(parent.begin(), parent.end(), 0);
    }

    int find(int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    void unite(int a, int b) {
        a = find(a);
        b = find(b);
        if (a == b) return;
        if (sz[a] < sz[b]) swap(a, b);
        parent[b] = a;
        sz[a] += sz[b];
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, K;
    cin >> N >> K;

    const int LIM = 1000000;
    vector<int> freq(LIM + 1, 0);

    int A = 0;
    for (int i = 0; i < N; i++) {
        int w;
        cin >> w;
        freq[w]++;
        A = max(A, w);
    }

    DSU dsu(A + 1);

    for (int d = K; d <= A; d++) {
        int first = -1;
        for (int m = d; m <= A; m += d) {
            if (freq[m] == 0) continue;
            if (first == -1) first = m;
            else dsu.unite(first, m);
        }
    }

    vector<long long> sum(A + 1, 0);
    long long ans = 0;

    for (int w = 1; w <= A; w++) {
        if (freq[w] == 0) continue;

        if (w < K) {
            ans = max(ans, (long long)w);
        } else {
            int r = dsu.find(w);
            sum[r] += 1LL * w * freq[w];
        }
    }

    for (long long x : sum) {
        ans = max(ans, x);
    }

    cout << ans << '\n';
    return 0;
}

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

投稿日時:
最終更新: