Problem p02985 (python)

Fixed submission:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
sys.setrecursionlimit(10 ** 6)


class FactorialMod:
    def __init__(self, prime, max_size):
        self.prime = prime
        self.max_size = max_size * 2
        self.fac = [0] * self.max_size
        self.finv = [0] * self.max_size
        self.factInit()

    def invMod(self, a):
        b, x, y = self.prime, 1, 0
        while b:
            t = a // b
            a -= t * b
            a, b = b, a
            x -= t * y
            x, y = y, x
        return x % self.prime

    def factInit(self):
        self.fac[0] = self.fac[1] = 1
        self.finv[0] = self.finv[1] = 1
        for i in range(2, self.max_size):
            self.fac[i] = (self.fac[i - 1] * i) % self.prime
            self.finv[i] = (self.finv[i - 1] * self.invMod(i)) % self.prime
    
    def fact(self, n):
        return self.fac[n]
    
    def factInv(self, n):
        return self.finv[n]
    

class Tree:
    def __init__(self, n):
        self.nodes = n
        self.parent = [-1 for _ in range(n + 1)]
        self.children = [[] for _ in range(n + 1)]
        self.edge = [[] for _ in range(n + 1)]
        self.depth = [0 for _ in range(n + 1)]
        self.root = -1
        self.max_depth = -1
    
    def addEdge(self, x, y):
        self.edge[x].append(y)
        self.edge[y].append(x)
    
    def prepare(self):
        checking = [i for i in range(self.nodes, 0, -1)]
        searching = [checking.pop()]
        next_level = []
        while searching:
            p = searching.pop()
            if self.root == -1:
                self.root = p
            cs = self.edge[p]
            self.children[p] = cs
            next_level.extend(cs)
            pd = self.depth[p]
            for c in cs:
                self.parent[c] = p
                self.depth[c] = pd + 1
                self.edge[c].remove(p)
            if searching == []:
                searching = next_level
        self.max_depth = max(self.depth)

    def getParent(self, x):
        return self.parent[x]
    
    def getRoot(self):
        if self.root != -1:
            return self.root
        for x in range(1, self.nodes + 1):
            if self.getParent(x) == -1:
                self.root = x
                return x
    
    def isRoot(self, x):
        return x == self.getRoot()

    def getChildren(self, x):
        return self.children[x]
    
    def getMaxDepth(self):
        return self.max_depth


def solve():
    PRIME = 10 ** 9 + 7
    N, K = list(map(int, input().split()))
    tree = Tree(N)
    f = FactorialMod(PRIME, K)
    for _ in range(N - 1):
        x, y = list(map(int, input().split()))
        tree.addEdge(x, y)
    tree.prepare()
    ans = 1
    searching = []
    next_searching = [[tree.getRoot()]]
    for i in range(tree.getMaxDepth() + 1):
        # print('Searching Level:', i)
        searching = next_searching[:]
        next_searching = []
        # print(searching)
        colors = K - min(2, i)  # 深さによって親と親の親がいる
        # print('colors:', colors)
        while searching:
            brothers = searching.pop()
            if colors < len(brothers):
                # print('Error: colors less than num of brothers')
                # print(colors, len(brothers))
                return 0
            ans *= f.fact(colors) * f.factInv(colors - len(brothers))
            ans %= PRIME
            # print('multiplied:', colors, 'P', len(brothers))
            for x in brothers:
                children = tree.getChildren(x)
                if children != []:
                    next_searching.append(children)
    return ans % PRIME


def main():
    print(solve())


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
sys.setrecursionlimit(10 ** 6)


class FactorialMod:
    def __init__(self, prime, max_size):
        self.prime = prime
        self.max_size = max_size * 2
        self.fac = [0] * self.max_size
        self.finv = [0] * self.max_size
        self.factInit()

    def invMod(self, a):
        b, x, y = self.prime, 1, 0
        while b:
            t = a // b
            a -= t * b
            a, b = b, a
            x -= t * y
            x, y = y, x
        return x % self.prime

    def factInit(self):
        self.fac[0] = self.fac[1] = 1
        self.finv[0] = self.finv[1] = 1
        for i in range(2, self.max_size):
            self.fac[i] = (self.fac[i - 1] * i) % self.prime
            self.finv[i] = (self.finv[i - 1] * self.invMod(i)) % self.prime
    
    def fact(self, n):
        return self.fac[n]
    
    def factInv(self, n):
        return self.finv[n]
    

class Tree:
    def __init__(self, n):
        self.nodes = n
        self.parent = [-1 for _ in range(n + 1)]
        self.children = [[] for _ in range(n + 1)]
        self.depth = [0 for _ in range(n + 1)]
        self.root = -1
        self.max_depth = -1
    
    def addEdge(self, x, y):
        if self.parent[y] == -1:
            self.parent[y] = x
            self.children[x].append(y)
        elif self.parent[x] == -1:
            self.parent[x] = y
            self.children[y].append(x)
        else:
            raise RuntimeError
    
    def calculateDepth(self):
        next_level = [self.getRoot()]
        d = 0
        while next_level:
            level = next_level[:]
            next_level = []
            while level:
                x = level.pop()
                self.depth[x] = d
                next_level += self.children[x]
            d += 1
        self.max_depth = d - 1
        
    def getParent(self, x):
        return self.parent[x]
    
    def getRoot(self):
        if self.root != -1:
            return self.root
        for x in range(1, self.nodes + 1):
            if self.getParent(x) == -1:
                self.root = x
                return x
    
    def isRoot(self, x):
        return x == self.getRoot()

    def getChildren(self, x):
        return self.children[x]
    
    def getMaxDepth(self):
        return self.max_depth


def solve():
    PRIME = 10 ** 9 + 7
    N, K = list(map(int, input().split()))
    tree = Tree(N)
    f = FactorialMod(PRIME, K)
    edges = []
    for _ in range(N - 1):
        x, y = list(map(int, input().split()))
        if x > y:
            x, y = y, x
        edges.append([x, y])
    edges = sorted(edges)
    for e in edges:
        tree.addEdge(e[0], e[1])
    tree.calculateDepth()
    ans = 1
    searching = []
    next_searching = [[tree.getRoot()]]
    for i in range(tree.getMaxDepth() + 1):
        # print('Searching Level:', i)
        searching = next_searching[:]
        next_searching = []
        # print(searching)
        colors = K - min(2, i)  # 深さによって親と親の親がいる
        # print('colors:', colors)
        while searching:
            brothers = searching.pop()
            if colors < len(brothers):
                # print('Error: colors less than num of brothers')
                # print(colors, len(brothers))
                return 0
            ans *= f.fact(colors) * f.factInv(colors - len(brothers))
            ans %= PRIME
            # print('multiplied:', colors, 'P', len(brothers))
            for x in brothers:
                children = tree.getChildren(x)
                if children != []:
                    next_searching.append(children)
    return ans % PRIME


def main():
    print(solve())


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
sys.setrecursionlimit(10 ** 6)


class FactorialMod:
    def __init__(self, prime, max_size):
        self.prime = prime
        self.max_size = max_size * 2
        self.fac = [0] * self.max_size
        self.finv = [0] * self.max_size
        self.factInit()

    def invMod(self, a):
        b, x, y = self.prime, 1, 0
        while b:
            t = a // b
            a -= t * b
            a, b = b, a
            x -= t * y
            x, y = y, x
        return x % self.prime

    def factInit(self):
        self.fac[0] = self.fac[1] = 1
        self.finv[0] = self.finv[1] = 1
        for i in range(2, self.max_size):
            self.fac[i] = (self.fac[i - 1] * i) % self.prime
            self.finv[i] = (self.finv[i - 1] * self.invMod(i)) % self.prime
    
    def fact(self, n):
        return self.fac[n]
    
    def factInv(self, n):
        return self.finv[n]
    

class Tree:
    def __init__(self, n):
        self.nodes = n
        self.parent = [-1 for _ in range(n + 1)]
        self.children = [[] for _ in range(n + 1)]
        self.depth = [0 for _ in range(n + 1)]
        self.root = -1
        self.max_depth = -1
    
    def addEdge(self, x, y):
        if self.parent[y] == -1:
            self.parent[y] = x
            self.children[x].append(y)
        elif self.parent[x] == -1:
            self.parent[x] = y
            self.children[y].append(x)
        else:
            raise RuntimeError
    
    def calculateDepth(self):
        next_level = [self.getRoot()]
        d = 0
        while next_level:
            level = next_level[:]
            next_level = []
            while level:
                x = level.pop()
                self.depth[x] = d
                next_level += self.children[x]
            d += 1
        self.max_depth = d - 1
        
    def getParent(self, x):
        return self.parent[x]
    
    def getRoot(self):
        if self.root != -1:
            return self.root
        for x in range(1, self.nodes + 1):
            if self.getParent(x) == -1:
                self.root = x
                return x
    
    def isRoot(self, x):
        return x == self.getRoot()

    def getChildren(self, x):
        return self.children[x]
    
    def getMaxDepth(self):
        return self.max_depth


def solve():
    PRIME = 10 ** 9 + 7
    N, K = list(map(int, input().split()))
    tree = Tree(N)
    f = FactorialMod(PRIME, K)
    for _ in range(N - 1):
        x, y = list(map(int, input().split()))
        tree.addEdge(x, y)
    tree.calculateDepth()
    ans = 1
    searching = []
    next_searching = [[tree.getRoot()]]
    for i in range(tree.getMaxDepth() + 1):
        # print('Searching Level:', i)
        searching = next_searching[:]
        next_searching = []
        # print(searching)
        colors = K - min(2, i)  # 深さによって親と親の親がいる
        # print('colors:', colors)
        while searching:
            brothers = searching.pop()
            if colors < len(brothers):
                # print('Error: colors less than num of brothers')
                # print(colors, len(brothers))
                return 0
            ans *= f.fact(colors) * f.factInv(colors - len(brothers))
            ans %= PRIME
            # print('multiplied:', colors, 'P', len(brothers))
            for x in brothers:
                children = tree.getChildren(x)
                if children != []:
                    next_searching.append(children)
    return ans % PRIME


def main():
    print(solve())


if __name__ == "__main__":
    main()