解説つき!競ってわかる はじめてpythonプログラミング

競プロの情報発信と備忘録を兼ねている

再帰関数を使用せずにUnion-Findをpythonで書いた

ABC157D-Friend Suggestions を解いてる最中に,Union-Findを使用したが,
PyPy3 で提出すると再帰関数が遅いことをうっかりしてTLEを出してしまった.

じゃあ,Python3で出せば良いじゃんと思うかもしれないが...

再帰の上限設定忘れ ->  RE(ちーーん)
となるのが,お決まりのパターン.

とまあ,こんな感じで,仮にpython3で提出するにしても,
sys.setrecursionlimit() の変更も必要なので,めんどくさい.

 

そこで,Union-Findの再帰関数部分を再帰関数にしなければ,TLEにはならないので,
変更しました(提出コード).

Union-Findがご入用の際には,ご活用くださいまし.

class UnionFind():
    """
    再帰関数を使用しないので,PyPy3で提出しやすい
    """

    def __init__(self, n):

        self.n = n
        self.parents = [-1]*n

    def unite(self, x, y):
        """unite y -> x"""

        px = self.find(x)
        py = self.find(y)
        if px == py:
            return 
        self.parents[py] = px
        return

    def find(self, x):

        dq = []
        while self.parents[x] != -1:
            dq.append(x)
            x = self.parents[x]
        else:
            while dq:
                y = dq.pop()
                self.parents[y] = x
        return x