Official

C - チームの旗の色 / Team Flag Colors Editorial by admin

GPT 5.4 High

概要

各操作は「2つのチームを合併する」なので、チーム管理には Union-Find(Disjoint Set Union, DSU)が最適です。
さらに、各チームの現在の旗の色と、「その色を持つチーム数」を辞書で管理すれば、最後の異なる色数をそのまま求められます。

考察

この問題では、選手ごとの所属チームが操作によってどんどん変わっていきます。
しかし操作の本質はずっと同じで、

  • \(U_i\) のいるチーム
  • \(V_i\) のいるチーム

をまとめて、できたチームの色を \(C_i\) にする、というものです。

重要な気づき 1: チーム管理は Union-Find でできる

最初は全員が別々のチームなので、各選手が 1 つの連結成分だと思えばよいです。
その後の操作は「2つの成分を統合する」なので、Union-Find がぴったりです。

Union-Find を使うと、

  • ある選手が今どのチームにいるか
  • 2人が同じチームかどうか
  • 2つのチームを合併する

を高速に処理できます。


重要な気づき 2: 各チームで知りたいのは「今の色」だけ

最終的に必要なのは、各チームの最終的な旗の色です。
チームのメンバー一覧そのものは必要ありません。

したがって、各 Union-Find の根に対して

  • そのチームの現在の色

だけを持てば十分です。

初期状態では色が未設定なので、コードでは -1 を「未設定」として使っています。


重要な気づき 3: 色の種類数は set だけでは足りない

最後に求めたいのは「現在存在するチームの旗の色として現れる異なる色の数」です。

一見すると set で色を管理したくなりますが、これは危険です。
なぜなら、同じ色を持つチームが複数存在できるからです。

例えば現在、

  • チーム A の色が \(5\)
  • チーム B の色も \(5\)

だったとします。
このとき set には {5} しか入りません。

もしチーム A が別の操作で消えたとき、単純に set から \(5\) を消してしまうと、
まだチーム B が色 \(5\) を持っているのに消えてしまって間違いです。

そこで必要なのが、

  • \(c\) を持つチームが今いくつあるか

を管理する辞書 color_count です。

例えば

  • color_count[5] = 2

なら、「色 5 のチームが 2 個ある」という意味です。
あるチームが消えたら 1 減らし、0 になったら辞書から削除します。

これで、現在存在する異なる色数 = len(color_count) になります。


素朴な方法がなぜ厳しいか

例えば、各操作のたびに

  • チームのメンバーを全部たどって合併する
  • 最後に全チームを作り直して色を集める

のような方法を考えると、合併のたびに大量の選手を動かすことになり、最悪で \(O(NM)\) 近くかかってしまいます。
\(N \le 200000,\ M \le 150000\) なので、これは間に合いません。

Union-Find なら、合併や代表元の取得をほぼ定数時間で処理できます。

アルゴリズム

以下を管理します。

  • parent[x]: Union-Find の親
  • size[x]: その集合の大きさ(union by size 用)
  • color[x]: 根 x が表すチームの現在の色(未設定なら -1
  • color_count[c]: 色 c を持つ現在のチーム数

補助関数

  • find(x): 選手 x が属するチームの根を返す
  • add_color(c): color_count[c] を 1 増やす
  • remove_color(c): color_count[c] を 1 減らし、0 なら削除する

各操作 \((u, v, c)\) の処理

1. 代表元を求める

ru = find(u)rv = find(v) を求めます。


2. すでに同じチームだった場合

ru == rv なら合併は起こりません。
ただし問題文通り、そのチームの旗の色は \(c\) に上書きされます。

処理は次の通りです。

  • そのチームに以前の色があれば、color_count から 1 つ減らす
  • 新しい色 \(c\) を設定する
  • color_count[c] を 1 増やす

3. 別々のチームだった場合

ru != rv なら、2つのチームは消えて新しい1チームになります。

ここで大事なのは、もとの2チームの色はもう存在しなくなることです。
だから先に、

  • ru の色が設定済みなら color_count から削除
  • rv の色が設定済みなら color_count から削除

を行います。

その後、Union-Find で 2 チームを合併します。
コードでは小さい集合を大きい集合にくっつける union by size を使っています。

最後に、合併後の根に色 \(c\) を設定し、

  • color[新しい根] = c
  • color_count[c] を 1 増やす

とします。


なぜこれで正しいか

各時点で color_count は、

「今存在している、色が設定済みのチームたちの色の個数」

を正しく表しています。

  • チームが色を上書きされたら、古い色を 1 つ消して新しい色を 1 つ追加
  • 2チームが合併したら、元の2チームの色を消して、新しい1チームの色を追加

としているためです。

したがって、すべての操作後に color_count のキーとして残っている色が、
最終的に存在するチームの色そのものです。

よって答えは len(color_count) です。

計算量

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

ここで \(\alpha(N)\) はアッカーマン関数の逆関数で、非常に小さい値です。
辞書の更新は平均 \(O(1)\) とみなせます。

実装のポイント

  • 入力は 1-indexed なので、コードでは u -= 1, v -= 1 して 0-indexed に直しています。

  • 色未設定を -1 で表しています。

  • 色の種類数は set ではなく、個数つき辞書で管理するのが重要です。

  • 吸収された側の根 rvcolor[rv] を消していなくても問題ありません。
    以後その頂点は根ではなくなるため、find を通して参照されるのは新しい根だけだからです。

    ソースコード

import sys

input = sys.stdin.readline

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

parent = list(range(N))
size = [1] * N
color = [-1] * N  # -1 means uncolored
color_count = {}


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


def add_color(c):
    color_count[c] = color_count.get(c, 0) + 1


def remove_color(c):
    v = color_count[c] - 1
    if v == 0:
        del color_count[c]
    else:
        color_count[c] = v


for _ in range(M):
    u, v, c = map(int, input().split())
    u -= 1
    v -= 1

    ru = find(u)
    rv = find(v)

    if ru == rv:
        if color[ru] != -1:
            remove_color(color[ru])
        color[ru] = c
        add_color(c)
    else:
        if color[ru] != -1:
            remove_color(color[ru])
        if color[rv] != -1:
            remove_color(color[rv])

        if size[ru] < size[rv]:
            ru, rv = rv, ru

        parent[rv] = ru
        size[ru] += size[rv]
        color[ru] = c
        add_color(c)

print(len(color_count))

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

posted:
last update: