算法分析

使用树形动态规划。对每个节点 u，维护两个状态：

dp[u][0]：选中 u 时，子树的最大权值和。
dp[u][d]（1 \leq d <k）：未选中 u 时，子树中最近选中节点距 u 为 d 的最大权值和。
dp[u][k]：未选中 u 且子树中无选中节点或最近选中节点距 u ≥k。

状态转移

选中 u：所有子节点 v 的最近选中节点距离必须 ≥k-1。
dp[u][0] = a_u + \sum_v \max_{d \geq k-1} dp[v][d]
未选中 u：对每个可能的父节点传递的距离 d_p，更新子节点状态：
- 若 d_p+1 <k，则 d_{child} = d_p+1。
- 否则，d_{child} =k。
  转移时，合并所有子节点在不同距离下的最大值。

复杂度

时间复杂度 O(nk)，空间复杂度 O(nk)。

代码实现

题外话：DeepSeek 在写代码的时候第十行末尾多打了一个右括号。以下代码没有这个问题。

import sys
from collections import defaultdict
sys.setrecursionlimit(1 << 25)

def main():
    n, k = map(int, sys.stdin.readline().split())
    a = list(map(int, sys.stdin.readline().split()))
    tree = [[] for _ in range(n)]
    for _ in range(n-1):
        u, v = map(int, sys.stdin.readline().split())
        tree[u-1].append(v-1)
        tree[v-1].append(u-1)

    # dp[u][d]: 在u的子树中，最近的选中节点距离u为d时的最大和
    dp = [defaultdict(int) for _ in range(n)]
    
    def dfs(u, parent):
        # 初始化，选中自己
        select = a[u]
        not_select = defaultdict(int)
        not_select[k] = 0  # 初始时，未选中且无选中节点
        
        for v in tree[u]:
            if v == parent:
                continue
            dfs(v, u)
            
            # 处理选中u的情况：子节点的最近选中距离 >=k-1
            max_sub = 0
            for d in dp[v]:
                if d >= k-1:
                    if dp[v][d] > max_sub:
                        max_sub = dp[v][d]
            select += max_sub
            
            # 处理未选中u的情况：合并子节点的状态
            new_not_select = defaultdict(int)
            for d_parent in not_select:
                current_sum = not_select[d_parent]
                max_v = 0
                for d_v in dp[v]:
                    new_d = d_parent +1 if d_parent +1 <k else k
                    if d_v == new_d:
                        candidate = dp[v][d_v]
                        if candidate > max_v:
                            max_v = candidate
                if max_v ==0:
                    # 子节点无法满足条件时，取所有可能的最大值
                    max_v = max(dp[v].values()) if dp[v] else 0
                current_sum += max_v
                key = min(d_parent +1, k)
                if key not in new_not_select or current_sum > new_not_select[key]:
                    new_not_select[key] = current_sum
            not_select = new_not_select
        
        # 更新dp[u]
        dp[u][0] = select
        for d in not_select:
            dp[u][d] = not_select[d]
    
    dfs(0, -1)
    
    # 根节点可以选中或未选中，取最大值
    max_val = 0
    for d in dp[0]:
        if d ==0 or d >=k:
            if dp[0][d] > max_val:
                max_val = dp[0][d]
    print(max_val)

if __name__ == "__main__":
    main()

【AI 提供】数据生成代码

import sys
import random
from random import randint

def generate_tree(n):
    edges = []
    for i in range(1, n):
        u = i
        v = randint(0, i-1)
        edges.append((u+1, v+1))  # 转换为1-based
    return edges

def main():
    n = randint(3, 10)
    k = randint(1, 3)
    a = [randint(1, 100) for _ in range(n)]
    edges = generate_tree(n)
    print(n, k)
    print(' '.join(map(str, a)))
    for u, v in edges:
        print(u, v)

if __name__ == "__main__":
    main()

【实际上】数据生成代码

import sys
from collections import defaultdict
import random
from random import randint
sys.setrecursionlimit(1 << 25)

def main(infile, outfile):
    n, k = map(int, infile.readline().split())
    a = list(map(int, infile.readline().split()))
    tree = [[] for _ in range(n)]
    for _ in range(n-1):
        u, v = map(int, infile.readline().split())
        tree[u-1].append(v-1)
        tree[v-1].append(u-1)

    # dp[u][d]: 在u的子树中，最近的选中节点距离u为d时的最大和
    dp = [defaultdict(int) for _ in range(n)]
    
    def dfs(u, parent):
        # 初始化，选中自己
        select = a[u]
        not_select = defaultdict(int)
        not_select[k] = 0  # 初始时，未选中且无选中节点
        
        for v in tree[u]:
            if v == parent:
                continue
            dfs(v, u)
            
            # 处理选中u的情况：子节点的最近选中距离 >=k-1
            max_sub = 0
            for d in dp[v]:
                if d >= k-1:
                    if dp[v][d] > max_sub:
                        max_sub = dp[v][d]
            select += max_sub
            
            # 处理未选中u的情况：合并子节点的状态
            new_not_select = defaultdict(int)
            for d_parent in not_select:
                current_sum = not_select[d_parent]
                max_v = 0
                for d_v in dp[v]:
                    new_d = d_parent +1 if d_parent +1 <k else k
                    if d_v == new_d:
                        candidate = dp[v][d_v]
                        if candidate > max_v:
                            max_v = candidate
                if max_v ==0:
                    # 子节点无法满足条件时，取所有可能的最大值
                    max_v = max(dp[v].values()) if dp[v] else 0
                current_sum += max_v
                key = min(d_parent +1, k)
                if key not in new_not_select or current_sum > new_not_select[key]:
                    new_not_select[key] = current_sum
            not_select = new_not_select
        
        # 更新dp[u]
        dp[u][0] = select
        for d in not_select:
            dp[u][d] = not_select[d]
    
    dfs(0, -1)
    
    # 根节点可以选中或未选中，取最大值
    max_val = 0
    for d in dp[0]:
        if d ==0 or d >=k:
            if dp[0][d] > max_val:
                max_val = dp[0][d]
    print(max_val, file=outfile)
    infile.close()
    outfile.close()

def generate_tree(n):
    edges = []
    for i in range(1, n):
        u = i
        v = randint(0, i-1)
        edges.append((u+1, v+1))  # 转换为1-based
    return edges

def gen(case):
    i = open(f'd{case}.in', 'a')
    k = randint(1, 3 if case == 1 else 10**3)
    n = randint(k, 10 if case == 1 else 10**5)
    a = [randint(1, 100) for _ in range(n)]
    edges = generate_tree(n)
    print(n, k, file=i)
    print(' '.join(map(str, a)), file=i)
    for u, v in edges:
        print(u, v, file=i)
    i.close()
    main(open(f'd{case}.in', 'r'), open(f'd{case}.out', 'a'))

if __name__ == "__main__":
    for case in range(1, 11):
        gen(case)

加载中... 加载中...