Cod sursa(job #2810965)

Utilizator LittleWhoFeraru Mihail LittleWho Data 30 noiembrie 2021 18:54:41
Problema Paduri de multimi disjuncte Scor 0
Compilator py Status done
Runda Arhiva educationala Marime 1.88 kb
class Forest:
    def __init__(self, n, roots=None, ranks=None):
        self.n = n
        
        if roots == None or ranks == None:
            self.roots = [i for i in range(n + 1)]
            self.ranks = [1 for i in range(n + 1)]
        else:
            self.root  = roots
            self.ranks = ranks

    def copy(self):
        return Forest(self.n, self.roots.copy(), self.ranks.copy())

    def in_same_set(self, x, y):
        assert x > 0 and y > 0
        
        return self.find(x) == self.find(y)

    def find(self, x):
        assert x > 0
        
        y = x
        
        while x != self.roots[x]:
            x = self.roots[x]
        
        while y != x:
            t = self.roots[y]
            self.roots[y] = x
            y = t
            
        return x

    def unite(self, x, y):
        assert x > 0 and y > 0
        
        x = self.find(x)
        y = self.find(y)

        # union by rank
        if self.ranks[x] > self.ranks[y]:
            self.roots[y] = x
            self.ranks[x] += self.ranks[y]
        else:
            self.roots[x] = y
            self.ranks[y] += self.ranks[x]


def tests():
    f = Forest(1000)
    f.unite(1, 2)
    f.unite(3, 4)
    assert f.in_same_set(1, 3) == False
    assert f.in_same_set(1, 2) == True
    f.unite(1, 3)
    assert f.in_same_set(1, 4) == True 
    
def infoarena():
    inf = [int(line.split(" ")) for line in open("disjoint.in").read().split("\n")]
    n, m = inf[0]
    
    f = Forest(n)
    output = []
    for op, a, b in inf[1:]:
        if op == 2:
            if f.in_same_set(a,b):
                output.append("DA")
            else:
                output.append("NU")
        else:
            f.unite(a, b)
                
    open("disjoint.out", "w").write('\n'.join(output))
    
if __name__ == "__main__":
    infoarena()