公式

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

GPT 5.4 High

Overview

Each operation “merges two teams,” so Union-Find (Disjoint Set Union, DSU) is the ideal data structure for team management.
Furthermore, by managing each team’s current flag color and “the number of teams with that color” using a dictionary, we can directly determine the number of distinct colors at the end.

Analysis

In this problem, each player’s team membership keeps changing through operations.
However, the essence of every operation is the same:

  • The team containing \(U_i\)
  • The team containing \(V_i\)

are merged, and the resulting team’s color is set to \(C_i\).

Key Insight 1: Team management can be done with Union-Find

Initially, everyone is on a separate team, so we can think of each player as their own connected component.
Subsequent operations “merge two components,” which is exactly what Union-Find is designed for.

Using Union-Find, we can efficiently handle:

  • Which team a player currently belongs to
  • Whether two players are on the same team
  • Merging two teams

Key Insight 2: All we need to know for each team is “the current color”

What we ultimately need is the final flag color of each team.
We don’t need the actual list of team members.

Therefore, it suffices to store only

  • the current color of that team

for each Union-Find root.

In the initial state, colors are unset, so the code uses -1 to represent “unset.”


Key Insight 3: A set alone is not enough for counting distinct colors

What we want to find at the end is “the number of distinct colors appearing as flag colors among currently existing teams.”

At first glance, you might want to manage colors with a set, but this is dangerous.
The reason is that multiple teams can have the same color.

For example, suppose currently:

  • Team A has color \(5\)
  • Team B also has color \(5\)

Then the set would only contain {5}.

If Team A disappears due to some operation and we simply remove \(5\) from the set,
Team B still has color \(5\), but it would be incorrectly removed.

What we need instead is a dictionary color_count that manages:

  • How many teams currently have color \(c\)

For example,

  • color_count[5] = 2

means “there are 2 teams with color 5.”
When a team disappears, we decrement by 1, and if it reaches 0, we delete it from the dictionary.

This way, the number of currently existing distinct colors = len(color_count).


Why a naive approach is too slow

For example, consider an approach where for each operation:

  • We traverse all members of the teams to merge them
  • At the end, we reconstruct all teams and collect colors

This would require moving a large number of players with each merge, resulting in worst-case close to \(O(NM)\) time.
Since \(N \le 200000,\ M \le 150000\), this is too slow.

With Union-Find, merging and finding representatives can be done in nearly constant time.

Algorithm

We maintain the following:

  • parent[x]: Union-Find parent
  • size[x]: Size of the set (for union by size)
  • color[x]: Current color of the team represented by root x (-1 if unset)
  • color_count[c]: Number of currently existing teams with color c

Helper Functions

  • find(x): Returns the root of the team that player x belongs to
  • add_color(c): Increments color_count[c] by 1
  • remove_color(c): Decrements color_count[c] by 1, and deletes the entry if it becomes 0

Processing each operation \((u, v, c)\)

1. Find the representatives

Compute ru = find(u) and rv = find(v).


2. If they are already on the same team

If ru == rv, no merge occurs.
However, as stated in the problem, the team’s flag color is overwritten to \(c\).

The processing is as follows:

  • If the team had a previous color, decrement it in color_count
  • Set the new color \(c\)
  • Increment color_count[c] by 1

3. If they are on different teams

If ru != rv, the two teams disappear and become one new team.

The important thing here is that the colors of the original two teams no longer exist.
So first:

  • If ru’s color is set, remove it from color_count
  • If rv’s color is set, remove it from color_count

Then, merge the two teams using Union-Find.
The code uses union by size, attaching the smaller set to the larger set.

Finally, set color \(c\) on the root after merging:

  • color[new root] = c
  • Increment color_count[c] by 1

Why this is correct

At each point in time, color_count correctly represents:

“The count of colors among currently existing teams that have a color set.”

This is because:

  • When a team’s color is overwritten, we remove one instance of the old color and add one instance of the new color
  • When two teams merge, we remove the colors of the original two teams and add the color of the new single team

Therefore, the colors remaining as keys in color_count after all operations
are exactly the colors of the teams that exist at the end.

Thus, the answer is len(color_count).

Complexity

  • Time complexity: \(O(M \alpha(N))\)
  • Space complexity: \(O(N + M)\)

Here, \(\alpha(N)\) is the inverse Ackermann function, which is an extremely small value.
Dictionary updates can be considered \(O(1)\) on average.

Implementation Notes

  • Input is 1-indexed, so the code converts to 0-indexed with u -= 1, v -= 1.

  • Unset color is represented by -1.

  • It is crucial to manage the number of distinct colors using a dictionary with counts, not a set.

  • It doesn’t matter if we don’t clear color[rv] for the absorbed root rv.
    Since that vertex is no longer a root, only the new root will be referenced through find.

    Source Code

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))

This editorial was generated by gpt-5.4-high.

投稿日時:
最終更新: