Official

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

gemini-3.5-flash-high

Overview

This problem asks us to find the maximum sum of specialty scores (total strength) within any group (connected component) when students satisfying a given condition (greatest common divisor is at least \(K\)) are grouped together.

We use Union-Find (DSU) to efficiently manage the connected components of the graph, and perform fast grouping by leveraging the relationships between divisors and multiples.


Analysis

1. Naive Approach and Its Limitations

The first idea that comes to mind is to determine whether \(\gcd(W_i, W_j) \ge K\) for all pairs of students \((i, j)\) and add an edge between pairs that satisfy the condition. However, since the number of students \(N\) is up to \(2 \times 10^5\), the number of pairs would be \(O(N^2) \approx 4 \times 10^{10}\), which will not run within the time limit (resulting in TLE).

2. Score Range and Reformulating the “Cooperation” Condition

Each student’s score \(W_i\) is at most \(M = 10^6\). We focus on the smallness of this value.

  • Students with \(W_i < K\) For any student \(j\), \(\gcd(W_i, W_j) \le W_i < K\), so they cannot cooperate with anyone. Thus, these students will always form a “group of size 1 containing only themselves”.
  • Students with \(W_i \ge K\) Students with the same score will always belong to the same group, because the greatest common divisor of their scores is their own score (\(\ge K\)).

Two students with different scores \(a, b \ge K\) can cooperate if and only if \(\gcd(a, b) = g \ge K\). This can be rephrased as: \(a\) and \(b\) are both multiples of some \(g \ge K\).” If both \(a\) and \(b\) are multiples of \(g\), then \(\gcd(a, b)\) is guaranteed to be a multiple of \(g\) (and thus at least \(g\)), so they can directly cooperate.

3. Fast Merging Using Harmonic Series

For each \(g \ge K\), we consider merging all students whose scores are multiples of \(g\) (\(g, 2g, 3g, \ldots\)) into the same group.

We iterate \(g\) from \(K\) to \(M\), and for each \(g\), we scan its multiples. The number of values we scan across all iterations is: $\( \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 \)\( By the properties of the "harmonic series," the total number of operations is approximately \)M \log M$, which is sufficiently fast.


Algorithm

  1. Handling \(W_i < K\) Since these students cannot form groups with anyone else, we record the maximum score among them beforehand (as an initial candidate for the maximum total strength).
  2. Calculating the Sum for Each Score For students with \(W_i \ge K\), we calculate the sum of scores for each distinct score value, storing it in sum_W[w].
  3. Initializing Union-Find Initialize a Union-Find data structure with elements from \(0\) to \(M\). The initial weight (total strength of the group) of each element \(w\) is set to sum_W[w].
  4. Merging Multiples Loop \(g\) from \(K\) to \(M\).
    • Iterate through the multiples of \(g\): \(x = g, 2g, 3g, \ldots \le M\).
    • Look for scores \(x\) that actually exist (i.e., sum_W[x] > 0).
    • Let the first such score found be first. Merge (Union) any subsequent existing multiples x into the same group as first. When merging, add up the weights of the roots in the Union-Find structure.
  5. Output the Answer Output the maximum of the total strengths of all groups in the Union-Find and the maximum score of a single student recorded in Step 1.

Complexity

  • Time Complexity: \(O(N + M \log M)\)

    • Processing \(W_i < K\) and constructing sum_W takes \(O(N)\) time.
    • The number of iterations in the loop for merging multiples is \(O(M \log M)\) by the harmonic series. Since each Union-Find operation (Find, Union) runs in nearly constant time \(O(\alpha(M))\), the merging part takes \(O(M \log M)\) time.
    • Overall, this easily runs within the time limit.
  • Space Complexity: \(O(N + M)\)

    • For the maximum score \(M \le 10^6\), the sizes of the parent and weight arrays in the Union-Find structure are \(O(M)\).

Implementation Points

  • Optimizing Union-Find (for Python) In Python, the find operation in Union-Find using recursive functions can be slow due to call overhead. In the provided code, the find operation is inlined using a while loop (written directly without function calls), and path compression is also performed inside the loop, achieving a significant speedup.

  • Skipping Unnecessary Merges By only targeting multiples where sum_W[x] > 0 (meaning students with that score actually exist) for merging, we avoid unnecessary Union-Find operations.

    Source Code

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()

This editorial was generated by gemini-3.5-flash-high.

posted:
last update: