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[新しい根] = ccolor_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ではなく、個数つき辞書で管理するのが重要です。吸収された側の根
rvのcolor[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: