Official

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


この問題は、次のように言い換えることができます。

\(K\le k\) なる整数 \(k\) すべてについて、専門スコアが \(k\) の倍数であるような \(2\) 人の学生の間すべてに辺を張る。 すべての連結成分にわたる、学生の専門スコアの合計の最大値を求めよ。

連想配列などを用いて、専門スコアが \(x\) である学生が存在すればそのような学生を \(1\) 人取得する(存在しなければそれを報告する)ようにしておきます。

事前に、同じ専門スコアの学生を適切に処理しておきます。 同じ専門スコアの学生が複数いた場合どのようにふるまうかは、学生たちの専門スコアと \(K\) との大小関係によって以下のように分類されます。

  • 学生たちの専門スコアが \(K\) 未満である場合、それらの学生はすべて単一の学生からなる連結成分になります。
  • 学生たちの専門スコアが \(K\) 以上である場合、それらの学生はすべて同一の連結成分に属します。

よって、専門スコアが \(K\) 以上で同一の値をもつ学生たちは、事前に辺を張っておくこととします。

以上の前処理によって、すべての整数 \(x\) に対して専門スコアが \(x\) である学生がたかだか \(1\) 人存在するとします。

専門スコア \(W _ i\) の最大値を \(M\) とします。\(K\le k\) なるすべての整数 \(k\) について、\(1\) 以上 \(M\) 以下の \(k\) の倍数をすべて走査することは \(O(M\log M)\) 時間で可能です。

あとは、それらの倍数のうちその専門スコアを持つ学生が存在するようなものに限って順に辺を張ることで言い換えたあとのグラフと同じ連結成分を得ることができます。

あとはそれぞれの連結成分に対して専門スコアの合計を求めることでこの問題を解くことができました。

実装例は以下のようになります。

#include <iostream>
#include <vector>
#include <unordered_map>
#include <atcoder/dsu>
using namespace std;

int main() {
    int N, K;
    cin >> N >> K;

    atcoder::dsu uf(N); // 連結成分を管理する
    unordered_map<int, int> score_to_index; // それぞれの専門スコアに対応する学生をひとり得る連想配列
    vector<int> W(N);
    int max_W = 0;
    for (int i = 0; i < N; ++i) {
        cin >> W[i];
        // すでに同じ専門スコアの学生がいて、スコアが K 以上なら
        if (score_to_index.contains(W[i]) && W[i] >= K) {
            uf.merge(i, score_to_index[W[i]]); // その学生との間に辺を張る
        }
        score_to_index[W[i]] = i;
        max_W = max(W[i], max_W);
    }

    // 各 k について
    for (int k = K; 2 * k <= max_W; ++k) {
        for (int w = k, prev = -1; w <= max_W; w += k) {
            if (score_to_index.contains(w)) { // k の倍数の学生がいれば
                if (prev != -1) {
                    uf.merge(score_to_index[w], prev); // 直前の k の倍数の学生と辺を張る
                }
                prev = score_to_index[w];
            }
        }
    }

    long ans = 0;
    for (vector group : uf.groups()) { // 連結成分ごとに
        long tmp = 0;
        for (int index : group) {
            tmp += W[index]; // 専門スコアの合計を求める
        }
        ans = max(ans, tmp); // 最大値が答え
    }
    cout << ans << endl;
    return 0;
}
from atcoder.dsu import DSU


N, K = map(int, input().split())

uf = DSU(N) # 連結成分を管理する
score_to_index = {} # それぞれの専門スコアに対応する学生をひとり得る連想配列
W = list(map(int, input().split()))
for i, w in enumerate(W):
    # すでに同じ専門スコアの学生がいて、スコアが K 以上なら
    if w in score_to_index and w >= K:
        uf.merge(i, score_to_index[w]) # その学生との間に辺を張る
    score_to_index[w] = i

max_W = max(W)

for k in range(K, max_W // 2 + 1):
    prev = -1
    for w in range(k, max_W + 1, k):
        if w in score_to_index: # k の倍数の学生がいれば
            if prev != -1:
                uf.merge(score_to_index[w], prev) # 直前の k の倍数の学生と辺を張る
            prev = score_to_index[w]

# 連結成分ごとの専門スコアの合計の最大値が答え
print(max(sum(W[index] for index in group) for group in uf.groups()))

posted:
last update: