kktsubota / compe Goto Github PK

class UnionFind(object):
    def __init__(self, n):
        self.parent = {i: i for i in range(n)}
        self.rank = {i: 0 for i in range(n)}
        self.size = {i: 1 for i in range(n)}
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])

        return self.parent[x]
    
    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return
        
        if self.rank[x] < self.rank[y]:
            self.parent[x] = y

        else:
            self.parent[y] = x

            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

        self.size[x] += self.size[y]
        self.size[y] = self.size[x]


    def get_size(self, x):
        return self.size[self.find(x)]

from collections import defaultdict


def prime_factorization(n):
    factors = defaultdict(int)
    i = 1
    n_sub = n
    while True:
        i += 1
        if i * i > n or n_sub == 1:
            break

        while n_sub % i == 0:
            n_sub = n_sub // i
            factors[i] += 1

    if n_sub != 1:
        factors[n_sub] = 1
    return factors

class ModInt(object):
    prime = 1000000007

    def __init__(self, x, calc_mod=True):
        if calc_mod:
            self.x = x % self.prime
        else:
            self.x = x
    
    def __add__(self, y):
        assert isinstance(y, ModInt)
        y = self.x + y.x
        if y.x >= self.prime:
            y.x -= self.prime
        return y
    
    def __sub__(self, y):
        assert isinstance(y, ModInt)
        y.x = self.x - y.x
        if y.x < 0:
            y.x += self.prime
        return y
    
    def __mul__(self, y):
        assert isinstance(y, ModInt)
        y = ModInt(self.x * y.x)
        return y
    
    def __pow__(self, y):
        if isinstance(y, ModInt):
            y = y.x
        assert isinstance(y, int) and y >= 0
        if y == 0:
            return ModInt(1, calc_mod=False)
        else:
            # y = 2 * z + {0,1}
            z = self.__pow__(y >> 1)
            z = z * z
            if (y & 1) == 0:
                return z
            else:
                return self * z

    def __truediv__(self, y):
        assert isinstance(y, ModInt)
        # Fermat's little theorem
        # y ** (p - 1) % p == 1
        # y ** (p - 2) % p == 1 / y
        return self * (y ** (self.prime - 2))
    
    def __floordiv__(self, y):
        raise NotImplementedError

    def __str__(self):
        return '{} (mod {})'.format(self.x, self.prime)