概要

\(N\) 人の選手に対して \(M\) 回のチーム合併＋旗色設定操作を行った後、旗の色が設定されているチームに存在する 異なる色の種類数 を求める問題です。Union-Find（素集合データ構造）を用いてチームの合併を効率的に管理します。

考察

重要な気づき

1. チームの合併 は Union-Find（Union by Rank + Path Compression）で効率的に処理できます。
2. 各操作では「2人の選手のチームを合併し、合併後のチームに色を設定する」という処理を行います。色は 上書き されるため、同じチームに複数回色が設定された場合、最後に設定された色だけが残ります。
3. 最終的に必要なのは「旗の色が設定されているチーム」の色の種類数です。一度も操作対象にならなかったチームは無視します。

素朴なアプローチの問題点

チームの合併をリストのコピーなどで管理すると、最悪 \(O(N)\) の時間がかかり、全体で \(O(NM)\) となって TLE の恐れがあります。Union-Find を使えば各操作をほぼ \(O(1)\)（正確には逆アッカーマン関数）で処理できます。

色の管理

Union-Find では各チームの 根（代表元） が一意に決まります。色は常に根に対して記録すれば、合併後も正しく管理できます。union(a, b) の結果として新しい根が返されるようにし、その根に色 \(C_i\) を設定します。

アルゴリズム

1. 初期化: 各選手 \(i\) を自分自身が親の Union-Find を用意する。色の配列 color[i] を None（未設定）で初期化する。

2. \(M\) 回の操作を順に処理:

   • 選手 \(U_i\) と \(V_i\) の所属するチームを union で合併する。
   • union は合併後の根（代表元）を返す。
   • その根に色 \(C_i\) を設定する（上書き）。

3. 最終結果の集計:

   • 全選手 \(1\) ～ \(N\) について find(i) で根を求め、重複なく全チームの根を列挙する。
   • 根に色が設定されている（color[root] is not None）チームの色を集合に追加する。
   • 集合のサイズが答え。

具体例

\(N=4, M=2\)、操作が \((1,2,\text{赤}), (3,4,\text{赤})\) の場合： - 操作1: 選手1と2が合併 → チーム{1,2} に色「赤」 - 操作2: 選手3と4が合併 → チーム{3,4} に色「赤」 - 最終的に色が設定されたチームは2つあるが、どちらも「赤」なので、答えは 1。

計算量

• 時間計算量: \(O((N + M) \cdot \alpha(N))\)
  • \(\alpha(N)\) は逆アッカーマン関数で実質定数。Union-Find の各操作（find, union）がほぼ \(O(1)\)。
  • 最後の集計で全選手を走査するのに \(O(N)\)。
• 空間計算量: \(O(N)\)
  • Union-Find の親配列・ランク配列・色配列にそれぞれ \(O(N)\)。

実装のポイント

• union の戻り値: 合併後の根を返すように実装することで、色を正しい根に設定できます。通常の Union-Find に少し手を加えるだけで実現できます。

• Path Compression（経路圧縮）: find の中で parent[x] = parent[parent[x]] とすることで木の高さを抑え、高速化しています。

• 最終集計時の重複排除: 同じ根を何度も数えないように seen_roots という集合を用いています。

• 色の上書き: 同じチームが複数回操作対象になると色が上書きされますが、常に根に色を記録しているため、最後の設定値が自然に残ります。

  ソースコード

import sys
input = sys.stdin.readline

def main():
    N, M = map(int, input().split())
    
    parent = list(range(N + 1))
    rank = [0] * (N + 1)
    color = [None] * (N + 1)  # color of the team represented by root
    
    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x
    
    def union(a, b):
        ra, rb = find(a), find(b)
        if ra == rb:
            return ra
        if rank[ra] < rank[rb]:
            ra, rb = rb, ra
        parent[rb] = ra
        if rank[ra] == rank[rb]:
            rank[ra] += 1
        return ra
    
    for _ in range(M):
        u, v, c = map(int, input().split())
        root = union(u, v)
        color[root] = c
    
    # Collect colors of all teams that have a color set
    seen_roots = set()
    colors = set()
    for i in range(1, N + 1):
        r = find(i)
        if r not in seen_roots:
            seen_roots.add(r)
            if color[r] is not None:
                colors.add(color[r])
    
    print(len(colors))

main()

この解説は claude4.6opus-thinking によって生成されました。