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. Я переписал этот код на плюсах, он зашёл, т.е. проблема не в алгоритме. Какие есть варианты почему не работает?


Ответы (0 шт):