runtime error при решении задачи на python
Задача:
дано дерево на n вершинах и число k, хотим найти наименьшее количество ребер, при удалении которых возникает хотя бы одна компонента размера
k(1 <= k <= n <= 1000).
Это вполне обычная задача на динамику, её решение даже можно нагуглить (1.5.2 https://archive.lksh.ru/2015/august/B/lectures/dp.ru.pdf)
вот мой код
def main():
n, k = map(int, input().split())
g = [[] for _ in range(n)] # g[v] - список соседей вершины v
for i in range(n - 1):
u, v = map(int, input().split()) # дерево задано списком ребер, ребро - пара вершин через пробел
u -= 1
v -= 1
g[v].append(u)
g[u].append(v)
INF = int(1e9 + 9)
dp = [[INF] * (n + 1) for i in range(n)]
sz = [1] * n # sz[v] - размер поддерева вершины v, если подвесить дерево за вершину 0
def dfs(v, par):
dp[v][1] = 0
for u in g[v]:
if u != par:
dfs(u, v)
newdp = [x + 1 for x in dp[v]]
for i in range(1, sz[v] + 1):
for j in range(1, sz[u] + 1):
newdp[i + j] = min(newdp[i + j], dp[v][i] + dp[u][j])
dp[v] = list(newdp)
sz[v] += sz[u]
#print(dp[v])
dfs(0, -1)
answer = dp[0][k]
for i in range(1, n):
answer = min(answer, dp[i][k] + 1)
print(answer)
main()
на нескольких тестах проверяющая система выдаёт вердикт RE. Я переписал этот код на плюсах, он зашёл, т.е. проблема не в алгоритме. Какие есть варианты почему не работает?