Problem p02985 (python)

Fixed submission:
from collections import defaultdict

mod = pow(10, 9) + 7

n, k = map(int, input().split())

graph = defaultdict(list)

for _ in range(n-1):
	a, b = map(int, input().split())
	a -= 1
	b -= 1
	graph[a].append(b)
	graph[b].append(a)

def dfs(v):
	global graph, k
	ans = k	# k choices of color for root
	stack = [v]
	vis_set = {v, }

	while stack:
		v = stack.pop()
		if v==0: 
			avail_count = k-1
		else: 
			avail_count = k-2
		for nxt in graph[v]:
			if nxt not in vis_set:
				if avail_count <= 0:
					return 0
				ans = (ans*avail_count)%mod
				avail_count -= 1
				stack.append(nxt)
				vis_set.add(nxt)
	return ans

ans = dfs(0)
print(ans)

Buggy submission (verdict: Runtime Error):
mod = pow(10, 9) + 7

n, k = map(int, input().split())

graph = {}

for _ in range(n-1):
	a, b = map(int, input().split())
	a -= 1
	b -= 1
	graph[a] = graph.get(a, []) + [b]
	graph[b] = graph.get(b, []) + [a]

def dfs(v):
	ans = k	# k choices of color for root
	stack = [v]
	vis_set = {v, }

	while stack:
		v = stack.pop()
		if v==0: avail_count = k-1
		else: avail_count = k-2
		for nxt in graph[v]:
			if nxt not in vis_set:
				if avail_count <= 0:
					return 0
				ans = (ans*avail_count)%mod
				avail_count -= 1
				stack.append(nxt)
				vis_set.add(nxt)
	return ans

ans = dfs(0)
print(ans)

Buggy submission (verdict: Runtime Error):
import math
from collections import defaultdict

mod = pow(10, 9) + 7

n, k = map(int, input().split())

graph = {}

for _ in range(n-1):
	a, b = map(int, input().split())
	a -= 1
	b -= 1
	graph[a] = graph.get(a, []) + [b]
	graph[b] = graph.get(b, []) + [a]

def dfs(v):
	ans = k	# k choices of color for root
	stack = [v]
	vis_set = {v, }

	while stack:
		v = stack.pop()
		if v==0: avail_count = k-1
		else: avail_count = k-2
		for nxt in graph[v]:
			if nxt not in vis_set:
				if avail_count <= 0:
					return 0
				ans = (ans*avail_count)%mod
				avail_count -= 1
				stack.append(nxt)
				vis_set.add(nxt)
	return ans

ans = dfs(0)
print(ans)

Buggy submission (verdict: Runtime Error):
import math

mod = pow(10, 9) + 7

n, k = map(int, input().split()) 

graph = {}

for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a] = graph.get(a, []) + [b]
	graph[b] = graph.get(b, []) + [a]

src = 1
E = []
vis_set = set()
depth = [0]*(n+1)
def dfs_eulerTour(src):
	global vis_set, depth
	vis_set.add(src)
	E.append(src)

	for neighbour in graph[src]:
		if not neighbour in vis_set:
			depth[neighbour] = depth[src] + 1
			dfs_eulerTour(neighbour)
			E.append(src)

dfs_eulerTour(src)
D = [depth[x] for x in E]
N = len(D)
def get_h():
	H = [-1]*(n+1)
	for i, ele in enumerate(E):
		if H[ele]==-1:
			vis_set.add(ele)
			H[ele] = i 
	return H
H = get_h()
M = [[0]*(int(math.log2(N))+1) for i in range(N+1)]

def preprocess(D, N, M):
	for i in range(N):
		M[i][0] = i

	for j in range(1, int(math.log2(N))+1):
		i = 0
		while i + (1 << j) - 1 < N:
			M[i][j] = M[i][j-1] if D[M[i][j-1]] <= D[M[i + (1 << (j-1))][j-1]] else M[i + (1 << (j-1))][j-1]
			i += 1
preprocess(D, N, M)

def get_dist(u, v):
	
	def RMQ(i, j):
		k = int(math.log2(j - i))
		return M[i][k] if D[M[i][k]] <= D[M[j - (1<<k)][k]] else M[j - (1<<k)][k]

	if H[u] == H[v]:
		return 0
	elif H[u]> H[v]:
		u, v = v, u
		q = RMQ(H[u], H[v])
		lca = E[q]
	else:
		q = RMQ(H[u], H[v])
		lca = E[q]
	return depth[u] + depth[v] - 2*depth[lca]

def free_verts(cur_vert, vis_set):
	count = sum(get_dist(cur_vert, other_vert)>2 for other_vert in vis_set)
	return count

def dfs(src):
	rem_color = k
	vis_set = {src,}
	stack = [src]
	color = [1]*(n+1)
	color[src] = k
	rem_color -= 1
	res = k
	while stack:
		v = stack.pop()
		for neighbour in graph[v]:
			if not neighbour in vis_set:
				vis_set.add(neighbour)
				stack.append(neighbour)
				add = free_verts(neighbour, vis_set)
				# print(neighbour, vis_set, add)
				color[neighbour] = (max(rem_color, 0) + add)%mod
				res = (res*color[neighbour])%mod
				rem_color -= 1
	return res
ans = dfs(src)
print(ans)

Buggy submission (verdict: Runtime Error):
import math
from collections import defaultdict

mod = pow(10, 9) + 7

n, k = map(int, input().split()) 

graph = defaultdict(list)

for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a].append(b)
	graph[b].append(a)

src = 1
E = []
vis_set = set()
depth = [0]*(n+1)
def dfs_eulerTour(src):
	global vis_set, depth
	vis_set.add(src)
	E.append(src)

	for neighbour in graph[src]:
		if not neighbour in vis_set:
			depth[neighbour] = depth[src] + 1
			dfs_eulerTour(neighbour)
			E.append(src)

dfs_eulerTour(src)
D = [depth[x] for x in E]
N = len(D)
def get_h():
	H = [-1]*(n+1)
	for i, ele in enumerate(E):
		if H[ele]==-1:
			vis_set.add(ele)
			H[ele] = i 
	return H
H = get_h()
M = [[0]*(int(math.log2(N))+1) for i in range(N+1)]

def preprocess(D, N, M):
	for i in range(N):
		M[i][0] = i

	for j in range(1, int(math.log2(N))+1):
		i = 0
		while i + (1 << j) - 1 < N:
			M[i][j] = M[i][j-1] if D[M[i][j-1]] <= D[M[i + (1 << (j-1))][j-1]] else M[i + (1 << (j-1))][j-1]
			i += 1
preprocess(D, N, M)

def get_dist(u, v):
	
	def RMQ(i, j):
		k = int(math.log2(j - i))
		return M[i][k] if D[M[i][k]] <= D[M[j - (1<<k)][k]] else M[j - (1<<k)][k]

	if H[u] == H[v]:
		return 0
	elif H[u]> H[v]:
		u, v = v, u
		q = RMQ(H[u], H[v])
		lca = E[q]
	else:
		q = RMQ(H[u], H[v])
		lca = E[q]
	return depth[u] + depth[v] - 2*depth[lca]

def free_verts(cur_vert, vis_set):
	count = sum(get_dist(cur_vert, other_vert)>2 for other_vert in vis_set)
	return count

def dfs(src):
	rem_color = k
	vis_set = {src,}
	stack = [src]
	color = [1]*(n+1)
	color[src] = k
	rem_color -= 1
	res = k
	while stack:
		v = stack.pop()
		for neighbour in graph[v]:
			if not neighbour in vis_set:
				vis_set.add(neighbour)
				stack.append(neighbour)
				add = free_verts(neighbour, vis_set)
				# print(neighbour, vis_set, add)
				color[neighbour] = (max(rem_color, 0) + add)%mod
				res = (res*color[neighbour])%mod
				rem_color -= 1
	return res
ans = dfs(src)
print(ans)

Buggy submission (verdict: Runtime Error):
import math
from collections import defaultdict
from functools import reduce

mod = pow(10, 9) + 7

n, k = map(int, input().split()) 

graph = defaultdict(list)

for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a].append(b)
	graph[b].append(a)

src = 1
E = []
vis_set = set()
depth = [0]*(n+1)
def dfs_eulerTour(src):
	global vis_set, depth
	vis_set.add(src)
	E.append(src)

	for neighbour in graph[src]:
		if not neighbour in vis_set:
			depth[neighbour] = depth[src] + 1
			dfs_eulerTour(neighbour)
			E.append(src)

dfs_eulerTour(src)
D = [depth[x] for x in E]
N = len(D)
def get_h():
	H = [-1]*(n+1)
	for i, ele in enumerate(E):
		if H[ele]==-1:
			vis_set.add(ele)
			H[ele] = i 
	return H
H = get_h()
M = [[0]*(int(math.log2(N))+1) for i in range(N+1)]

def preprocess(D, N, M):
	for i in range(N):
		M[i][0] = i

	for j in range(1, int(math.log2(N))+1):
		i = 0
		while i + (1 << j) - 1 < N:
			M[i][j] = M[i][j-1] if D[M[i][j-1]] <= D[M[i + (1 << (j-1))][j-1]] else M[i + (1 << (j-1))][j-1]
			i += 1
preprocess(D, N, M)

def get_dist(u, v):
	
	def RMQ(i, j):
		k = int(math.log2(j - i))
		return M[i][k] if D[M[i][k]] <= D[M[j - (1<<k)][k]] else M[j - (1<<k)][k]

	if H[u] == H[v]:
		return 0
	elif H[u]> H[v]:
		u, v = v, u
		q = RMQ(H[u], H[v])
		lca = E[q]
	else:
		q = RMQ(H[u], H[v])
		lca = E[q]
	return depth[u] + depth[v] - 2*depth[lca]

def free_verts(cur_vert, vis_set):
	count = sum(get_dist(cur_vert, other_vert)>2 for other_vert in vis_set)
	return count

def dfs(src):
	rem_color = k
	vis_set = {src,}
	stack = [src]
	color = [1]*(n+1)
	color[src] = k
	rem_color -= 1
	while stack:
		v = stack.pop()
		for neighbour in graph[v]:
			if not neighbour in vis_set:
				vis_set.add(neighbour)
				stack.append(neighbour)
				add = free_verts(neighbour, vis_set)
				# print(neighbour, vis_set, add)
				color[neighbour] = (max(rem_color, 0) + add)%mod
				rem_color -= 1
	return reduce(lambda x, y: (x*y)%mod, color)
ans = dfs(src)
print(ans)

Buggy submission (verdict: Runtime Error):
import math
from collections import defaultdict
from functools import reduce

mod = pow(10, 9) + 7

n, k = map(int, input().split()) 

graph = defaultdict(list)

for _ in range(n-1):
	a, b = map(int, input().split())
	graph[a].append(b)
	graph[b].append(a)

src = 1
E = []
vis_set = set()
depth = [0]*(n+1)
def dfs_eulerTour(src):
	global vis_set, depth
	vis_set.add(src)
	E.append(src)

	for neighbour in graph[src]:
		if not neighbour in vis_set:
			depth[neighbour] = depth[src] + 1
			dfs_eulerTour(neighbour)
			E.append(src)

dfs_eulerTour(src)
D = [depth[x] for x in E]
N = len(D)
def get_h():
	H = [-1]*(n+1)
	for i, ele in enumerate(E):
		if H[ele]==-1:
			vis_set.add(ele)
			H[ele] = i 
	return H
H = get_h()
M = [[0]*(int(math.log2(N))+1) for i in range(N+1)]

def preprocess(D, N, M):
	for i in range(N):
		M[i][0] = i

	for j in range(1, int(math.log2(N))+1):
		i = 0
		while i + (1 << j) - 1 < N:
			M[i][j] = M[i][j-1] if D[M[i][j-1]] <= D[M[i + (1 << (j-1))][j-1]] else M[i + (1 << (j-1))][j-1]
			i += 1
preprocess(D, N, M)

def get_dist(u, v):
	
	def RMQ(i, j):
		k = int(math.log2(j - i))
		return M[i][k] if D[M[i][k]] <= D[M[j - (1<<k)][k]] else M[j - (1<<k)][k]

	if H[u] == H[v]:
		return 0
	elif H[u]> H[v]:
		u, v = v, u
		q = RMQ(H[u], H[v])
		lca = E[q]
	else:
		q = RMQ(H[u], H[v])
		lca = E[q]
	return depth[u] + depth[v] - 2*depth[lca]

def free_verts(cur_vert, vis_set):
	count = sum(get_dist(cur_vert, other_vert)>2 for other_vert in vis_set)
	return count

def dfs(src):
	rem_color = k
	vis_set = {src,}
	stack = [src]
	color = [1]*(n+1)
	color[src] = k
	rem_color -= 1
	while stack:
		v = stack.pop()
		for neighbour in graph[v]:
			if not neighbour in vis_set:
				vis_set.add(neighbour)
				stack.append(neighbour)
				add = free_verts(neighbour, vis_set)
				# print(neighbour, vis_set, add)
				color[neighbour] = (max(rem_color, 0) + add)%mod
				rem_color -= 1
	return reduce(lambda x, y: (x*y)%mod, color)
ans = dfs(src)
print(ans)