Get the size (number of elements) of UnionFind in Python
I'm sorry. Second notebook for myself
Last time I wrote the minimum implementation of Union-Find, and now UnionFind is unmatched! I thought
ABC157 D-Friend Suggestion
https://atcoder.jp/contests/abc157/tasks/abc157_d
AtCoder's problem has made me confused.
The cause is that the implementation to get the size (number of elements) has not been done. Let's reimplement.
Implementation code
N,Q = map(int,input().split())
par = [-1]*(N+1)
def find(x):
if par[x] < 0:
return x
else:
par[x] = find(par[x]) #Path compression
return par[x]
def same(x,y):
return find(x) == find(y)
def unite(x,y):
x = find(x)
y = find(y)
if x == y:
return 0
else:
if par[x] > par[y]:
x,y = y,x
par[x] += par[y]
par[y] = x
def size(x):
return -par[find(x)]
This is the revised version.
Initialization
par = [-1]*(N+1)
It should be noted that all the initial initial values are set to -1 and managed.
If you are a parent, you are judged by whether you are negative or not.
Unite
def unite(x,y):
x = find(x)
y = find(y)
if x == y:
return 0
else:
if par[x] > par[y]:
x,y = y,x
par[x] += par[y]
par[y] = x
Is to take the negative sum of both parents when joining.
By doing so, the child element of par will contain the number of the parent node.
The size (number of elements) of the tree can be held as a negative value in the parent element of par.
Example
For example, in this image, it will be [-5, 0, 0, 0, 0].
size
Now it's clear why the return value of the size function has a minus ‘ー’!
def size(x):
return -par[find(x)]
Summary
I was surprised at first that a minus appeared before the return even though I only got the size, but it was very easy to understand once I understood it.
This time, Union Find should be able to match!
Recommended Posts