/ SeriousOJ /

Record Detail

Wrong Answer


  
# Status Time Cost Memory Cost
#1 Wrong Answer 13ms 2.945 MiB
#2 Wrong Answer 14ms 3.105 MiB

Code

import sys
sys.setrecursionlimit(10**7)

n, K = map(int, input().split())
values = list(map(int, input().split()))

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

# dp[node][k] = max path sum of length k ending at node
dp = [[float('-inf')] * (K+1) for _ in range(n+1)]
visited = [False] * (n+1)

# Initialize dp for path length 1 (just the node itself)
for i in range(1, n+1):
    dp[i][1] = values[i-1]

ans = 0

def dfs(node, parent):
    global ans
    visited[node] = True
    
    for child in graph[node]:
        if child != parent:
            dfs(child, node)
            
            # Update dp[node] with child's dp
            for k_len in range(2, K+1):
                dp[node][k_len] = max(dp[node][k_len], dp[child][k_len-1] + values[node-1])

    # Update with parent's dp
    if parent != -1:
        for k_len in range(2, K+1):
            dp[node][k_len] = max(dp[node][k_len], dp[parent][k_len-1] + values[node-1])

    ans = max(ans, max(dp[node]))

# Run DFS from root (node 1)
dfs(1, -1)

# Print the result (if no valid path found, dp will remain -inf)
print(ans if ans != float('-inf') else 0)

Information

Submit By
Type
Submission
Problem
P1161 Max path sum (Hard Version)
Contest
Brain Booster #8
Language
Python 3 (Python 3.12.3)
Submit At
2025-02-17 16:43:11
Judged At
2025-02-17 16:43:11
Judged By
Score
0
Total Time
14ms
Peak Memory
3.105 MiB