C - チームの旗の色 / Team Flag Colors 解説 by admin
GPT 5.4 HighOverview
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 parentsize[x]: Size of the set (for union by size)color[x]: Current color of the team represented by rootx(-1if unset)color_count[c]: Number of currently existing teams with colorc
Helper Functions
find(x): Returns the root of the team that playerxbelongs toadd_color(c): Incrementscolor_count[c]by 1remove_color(c): Decrementscolor_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 fromcolor_count - If
rv’s color is set, remove it fromcolor_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 rootrv.
Since that vertex is no longer a root, only the new root will be referenced throughfind.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.
投稿日時:
最終更新: