Connectivity Queries (DSU)
Programming Interview
Hard
4 views
Problem Description
Consider {x}. Make DSU class with find and union. Queries: UNION u v, ASK u v. For ASK output YES if connected else NO.
Input Format
First line n q. Next q lines queries.
Output Format
Outputs for ASK lines.
Sample Test Case
Input:
5 6
UNION 1 2
ASK 1 3
UNION 2 3
ASK 1 3
ASK 4 5
UNION 4 5
Official Solution
import sys
lines=sys.stdin.read().splitlines()
if not lines: sys.exit(0)
first=lines[0].split()
if len(first)<2: sys.exit(0)
n=int(first[0]); q=int(first[1])
class DSU:
def __init__(self,n):
self.p=list(range(n+1))
self.sz=[1]*(n+1)
def find(self,x):
while self.p[x]!=x:
self.p[x]=self.p[self.p[x]]
x=self.p[x]
return x
def union(self,a,b):
ra=self.find(a)
rb=self.find(b)
if ra==rb:
return
if self.sz[ra]<self.sz[rb]:
ra,rb=rb,ra
self.p[rb]=ra
self.sz[ra]+=self.sz[rb]
dsu=DSU(n)
out=[]
for i in range(1,1+q):
parts=(lines[i] if i<len(lines) else '').split()
if len(parts)<3: continue
u=int(parts[1]); v=int(parts[2])
if parts[0]=='UNION':
dsu.union(u,v)
else:
out.append('YES' if dsu.find(u)==dsu.find(v) else 'NO')
sys.stdout.write('\
'.join(out))
Solutions (0)
No solutions submitted yet. Be the first!
No comments yet. Start the discussion!