[HackerRank][Python3] Friend Circle Queries

Problem :

https://www.hackerrank.com/challenges/friend-circle-queries/problem

My Solution :

#!/usr/bin/env python3


class Node(object):
    def __init__(self, data):
        self.data = data
        self.parent = self
        self.rank = 1
        self.count = 1


def find_parent(node):
    if node.parent is node:
        return node
    return find_parent(node.parent)


def union_nodes(n1, n2):
    if n1 is n2:
        return n1
    p1, p2 = find_parent(n1), find_parent(n2)
    if p1 is p2:
        return p1
    if p1.rank < p2.rank:
        p1, p2 = p2, p1
    p2.parent = p1
    if p1.rank == p2.rank:
        p1.rank += 1
    p1.count += p2.count
    return p1


def maxCircle(queries):
    parents = {}
    M = 0
    for d1, d2 in queries:
        p1 = parents.get(d1) or Node(d1)
        p2 = parents.get(d2) or Node(d2)
        p = union_nodes(p1, p2)
        parents[d1] = p
        parents[d2] = p
        M = max(M, p.count)
        yield M


q = int(input())
queries = []
for _ in range(q):
    queries.append(list(map(int, input().rstrip().split())))
for ans in maxCircle(queries):
    print(ans)

Comment :

disjoint-set 문제를 나름대로 최적화 해서 구현해 보았다. 풀고 나서 좀 뿌듯했다. 역시 트리가 좋구나.