Official

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

Gemini 3.0 Flash (Thinking)

Overview

This is a problem where you merge teams that \(N\) players belong to while overwriting the “flag color” for each team. You need to determine the flag color held by each team that exists at the end, and find the number of distinct colors.

Analysis

The key point of this problem is how to efficiently manage “team merging” and “color updates”.

  1. Team Merging (Using Union-Find) The operation “merge the teams that players \(U_i\) and \(V_i\) belong to” can be processed very efficiently using the Union-Find (Disjoint Set Union) data structure. With Union-Find, you can manage which group (team) each player belongs to through a “representative (root)”.

  2. Color Management Each team has one color. Since the color is overwritten every time a team merger occurs, it is efficient to associate the color with the current representative of the team.

    • When a merger occurs: Set color \(C_i\) on the new team’s representative.
    • When they are already on the same team: Update the current representative’s color to \(C_i\).
  3. Why doesn’t a simple array work? If you rewrite “which color each player has” for all members, a single operation takes up to \(O(N)\), resulting in \(O(NM)\) overall, which is too slow. By using Union-Find and having only the representative hold the color information, significant speedup is possible.

Algorithm

The solution is computed with the following steps.

  1. Initialization:
    • Create a Union-Find tree with \(N\) elements.
    • Prepare an array color to hold each team’s color, initialized to -1 meaning “not set”.
  2. Processing Operations: For each of the \(M\) operations \((U_i, V_i, C_i)\), do the following:
    • Find the representatives \(root\_u, root\_v\) of \(U_i\) and \(V_i\).
    • If the two are on different teams:
      • Merge the two trees and determine the new representative (Union by Rank).
      • Update the new representative’s color to \(C_i\).
    • If the two are already on the same team:
      • Update that representative’s color to \(C_i\).
  3. Aggregation:
    • Check all players \(i = 1 \dots N\).
    • If player \(i\) is the representative of their team (parent[i] == i) and color[i] is not -1, add that color to a set (set).
    • Finally, output the number of elements in the set.

Complexity

  • Time Complexity: \(O(M \alpha(N) + N)\)
    • The \(M\) Union-Find operations take \(O(M \alpha(N))\) (\(\alpha\) is the inverse Ackermann function, which is practically constant).
    • The final aggregation takes \(O(N)\).
    • Overall, it runs very efficiently.
  • Space Complexity: \(O(N)\)
    • The Union-Find arrays (parent, rank) and the color management array use \(O(N)\) memory.

Implementation Notes

  • Iterative Union-Find: Since Python has a recursion depth limit, the find operation (finding the representative) is implemented iteratively using a while loop to prevent runtime errors.

  • Path Compression: During the find operation, parents are directly linked to the representative, speeding up subsequent operations.

  • Fast I/O: Since \(N, M\) can be large, input is read all at once using sys.stdin.read().split() for efficiency.

    Source Code

import sys

def solve():
    # 標準入力からすべてのデータを読み込み、スペースで分割してリストにする
    # 大量の入力を処理する場合、この方法が一般的に高速です
    try:
        input_data = sys.stdin.read().split()
    except EOFError:
        return
        
    if not input_data:
        return
    
    # 整数に変換
    data = list(map(int, input_data))
    
    N = data[0]
    M = data[1]
    
    # Union-Find用のデータ構造
    # parent[i]: 選手 i が属するチームの親。自分自身ならそのチームの代表者
    # rank[i]: 木の高さ(Union by Rank用)
    # color[i]: 代表者 i が率いるチームの現在の旗の色。-1は未設定を表す
    parent = list(range(N + 1))
    rank = [0] * (N + 1)
    color = [-1] * (N + 1)
    
    data_idx = 2
    for _ in range(M):
        # 各操作の情報を取得
        if data_idx + 2 >= len(data):
            break
        u = data[data_idx]
        v = data[data_idx+1]
        c = data[data_idx+2]
        data_idx += 3
        
        # 選手 u が属するチームの代表者を見つける(経路圧縮付き、反復型)
        root_u = u
        while parent[root_u] != root_u:
            root_u = parent[root_u]
        curr = u
        while curr != root_u:
            parent[curr], curr = root_u, parent[curr]
            
        # 選手 v が属するチームの代表者を見つける(経路圧縮付き、反復型)
        root_v = v
        while parent[root_v] != root_v:
            root_v = parent[root_v]
        curr = v
        while curr != root_v:
            parent[curr], curr = root_v, parent[curr]
        
        if root_u != root_v:
            # 異なるチームを合併する(Union by Rank)
            if rank[root_u] < rank[root_v]:
                root_u, root_v = root_v, root_u
            parent[root_v] = root_u
            if rank[root_u] == rank[root_v]:
                rank[root_u] += 1
            # 合併後のチーム(新しい代表者 root_u)に新しい色を設定
            # 以前の root_v の色は、root_v が代表者でなくなったため無視される
            color[root_u] = c
        else:
            # すでに同じチームに属している場合、そのチームの色を更新
            color[root_u] = c
            
    # すべての操作終了後、各チーム(代表者)の旗の色を確認
    unique_colors = set()
    for i in range(1, N + 1):
        # parent[i] == i ならば、その選手はチームの代表者
        if parent[i] == i:
            # color[i] が -1 でなければ、そのチームには旗の色が設定されている
            if color[i] != -1:
                unique_colors.add(color[i])
    
    # 異なる色の種類数を出力
    sys.stdout.write(str(len(unique_colors)) + '\n')

if __name__ == '__main__':
    solve()

This editorial was generated by gemini-3-flash-thinking.

posted:
last update: