Helga Hufflepuff's Cup (Codeforces 855C)
Tags: DP, trees

Harry, Ron and Hermione have figured out that Helga Hufflepuff's cup is a horcrux. Through her encounter with Bellatrix Lestrange, Hermione came to know that the cup is present in Bellatrix's family vault in Gringott's Wizarding Bank. The Wizarding bank is in the form of a tree with total n vaults where each vault has some type, denoted by a number between 1 to m. A tree is an undirected connected graph with no cycles.The vaults with the highest security are of type k, and all vaults of type k have the highest security.There can be at most x vaults of highest security. Also, if a vault is of the highest security, its adjacent vaults are guaranteed to not be of the highest security and their type is guaranteed to be less than k.Harry wants to consider every possibility so that he can easily find the best path to reach Bellatrix's vault. So, you have to tell him, given the tree structure of Gringotts, the number of possible ways of giving each vault a type such that the above conditions hold.

Input: The first line of input contains two space separated integers, n and m — the number of vaults and the number of different vault types possible. (1 ≤ n ≤ 105, 1 ≤ m ≤ 109).Each of the next n - 1 lines contain two space separated integers ui and vi (1 ≤ ui, vi ≤ n) representing the i-th edge, which shows there is a path between the two vaults ui and vi. It is guaranteed that the given graph is a tree.The last line of input contains two integers k and x (1 ≤ k ≤ m, 1 ≤ x ≤ 10), the type of the highest security vault and the maximum possible number of vaults of highest security.

Output: Output a single integer, the number of ways of giving each vault a type following the conditions modulo 109 + 7.

Note: In test case 1, we cannot have any vault of the highest security as its type is 1 implying that its adjacent vaults would have to have a vault type less than 1, which is not allowed. Thus, there is only one possible combination, in which all the vaults have type 2.

Sample Inputs:
4 2
1 2
2 3
1 4
1 2

3 3
1 2
1 3
2 1

3 1
1 2
1 3
1 1

Sample Outputs:
1

13

0

Solution (Python):
import sys
from array import array
from collections import deque

class graph:

    def __init__(self, n):
        (self.n, self.gdict) = (n, [array('i') for _ in range(n + 1)])

    def add_edge(self, node1, node2):
        self.gdict[node1].append(node2)
        self.gdict[node2].append(node1)

    def subtree(self, root):
        (queue, visit, child) = (deque([root]), array('b', [False] * (self.n + 1)), array('i'))
        dp = [[array('i', [0] * (self.n + 1)) for _ in range(max_ + 1)] for _ in range(3)]
        (par, visit[root]) = (array('i', [-1] * (self.n + 1)), True)
        while queue:
            s = queue.popleft()
            for i1 in self.gdict[s]:
                if not visit[i1]:
                    queue.append(i1)
                    visit[i1] = True
                    child.append(i1)
                    par[i1] = s
            dp[0][0][s] = high - 1
            dp[2][0][s] = m - high
            dp[1][1][s] = 1
        for i in child[::-1]:
            p = par[i]
            dp1 = [0] * (max_ + 1)
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(add(dp[0][j1][i], dp[2][j1][i]), dp[1][j1][i])
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[0][j][p], su))
            for j in range(max_ + 1):
                dp[0][j][p] = dp1[j]
            dp1 = [0] * (max_ + 1)
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(dp[2][j1][i], dp[0][j1][i])
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[2][j][p], su))
            for j in range(max_ + 1):
                dp[2][j][p] = dp1[j]
            dp1 = [0] * (max_ + 1)
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[1][j][p], dp[0][j1][i]))
            for j in range(1, max_ + 1):
                dp[1][j][p] = dp1[j]
        ans = 0
        for i in range(3):
            for j in range(max_ + 1):
                ans = add(ans, dp[i][j][1])
        print(ans)
input = lambda : sys.stdin.buffer.readline().decode().strip()
mult = lambda a, b: a * b % mod
add = lambda a, b: (a + b) % mod
(n, m) = map(int, input().split())
(g, mod) = (graph(n), 10 ** 9 + 7)
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    g.add_edge(u, v)
(high, max_) = map(int, input().split())
g.subtree(1)

Solution (Python):
import sys
from array import array
from collections import deque

class graph:

    def __init__(self, n):
        (self.n, self.gdict) = (n, [array('i') for _ in range(n + 1)])

    def add_edge(self, node1, node2):
        self.gdict[node1].append(node2)
        self.gdict[node2].append(node1)

    def subtree(self, root):
        (queue, visit, child) = (deque([root]), array('b', [False] * (self.n + 1)), array('i'))
        dp = [[array('i', [0] * (self.n + 1)) for _ in range(max_ + 1)] for _ in range(3)]
        (par, visit[root]) = (array('i', [-1] * (self.n + 1)), True)
        dp1 = [array('i', [0] * (max_ + 1)) for _ in range(3)]
        while queue:
            s = queue.popleft()
            for i1 in self.gdict[s]:
                if not visit[i1]:
                    queue.append(i1)
                    visit[i1] = True
                    child.append(i1)
                    par[i1] = s
            dp[0][0][s] = high - 1
            dp[2][0][s] = m - high
            dp[1][1][s] = 1
        for i in child[::-1]:
            p = par[i]
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(add(dp[0][j1][i], dp[2][j1][i]), dp[1][j1][i])
                    dp1[0][j + j1] = add(dp1[0][j + j1], mult(dp[0][j][p], su))
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(dp[2][j1][i], dp[0][j1][i])
                    dp1[2][j + j1] = add(dp1[2][j + j1], mult(dp[2][j][p], su))
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    dp1[1][j + j1] = add(dp1[1][j + j1], mult(dp[1][j][p], dp[0][j1][i]))
            for _ in range(3):
                for j in range(max_ + 1):
                    dp[_][j][p] = dp1[_][j]
                    dp1[_][j] = 0
        ans = 0
        for i in range(3):
            for j in range(max_ + 1):
                ans = add(ans, dp[i][j][1])
        print(ans)
input = lambda : sys.stdin.buffer.readline().decode().strip()
mult = lambda a, b: a * b % mod
add = lambda a, b: (a + b) % mod
(n, m) = map(int, input().split())
(g, mod) = (graph(n), 10 ** 9 + 7)
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    g.add_edge(u, v)
(high, max_) = map(int, input().split())
g.subtree(1)

Solution (Python):
import sys
from array import array
from collections import deque

class graph:

    def __init__(self, n):
        (self.n, self.gdict) = (n, [array('i') for _ in range(n + 1)])

    def add_edge(self, node1, node2):
        self.gdict[node1].append(node2)
        self.gdict[node2].append(node1)

    def subtree(self, root):
        (queue, visit, child) = (deque([root]), array('b', [False] * (self.n + 1)), array('i'))
        dp = [[array('i', [0] * (self.n + 1)) for _ in range(max_ + 1)] for _ in range(3)]
        (par, visit[root]) = (array('i', [-1] * (self.n + 1)), True)
        while queue:
            s = queue.popleft()
            for i1 in self.gdict[s]:
                if not visit[i1]:
                    queue.append(i1)
                    visit[i1] = True
                    child.append(i1)
                    par[i1] = s
            dp[0][0][s] = high - 1
            dp[2][0][s] = m - high
            dp[1][1][s] = 1
        for i in child[::-1]:
            p = par[i]
            dp1 = array('i', [0] * (max_ + 1))
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(add(dp[0][j1][i], dp[2][j1][i]), dp[1][j1][i])
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[0][j][p], su))
            for j in range(max_ + 1):
                dp[0][j][p] = dp1[j]
            dp1 = array('i', [0] * (max_ + 1))
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    su = add(dp[2][j1][i], dp[0][j1][i])
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[2][j][p], su))
            for j in range(max_ + 1):
                dp[2][j][p] = dp1[j]
            dp1 = array('i', [0] * (max_ + 1))
            for j in range(max_ + 1):
                for j1 in range(max_ - j + 1):
                    dp1[j + j1] = add(dp1[j + j1], mult(dp[1][j][p], dp[0][j1][i]))
            for j in range(1, max_ + 1):
                dp[1][j][p] = dp1[j]
        ans = 0
        for i in range(3):
            for j in range(max_ + 1):
                ans = add(ans, dp[i][j][1])
        print(ans)
input = lambda : sys.stdin.buffer.readline().decode().strip()
mult = lambda a, b: a * b % mod
add = lambda a, b: (a + b) % mod
(n, m) = map(int, input().split())
(g, mod) = (graph(n), 10 ** 9 + 7)
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    g.add_edge(u, v)
(high, max_) = map(int, input().split())
g.subtree(1)