Bootstrap

【用deepseek和chatgpt做算法竞赛】——华为算法精英实战营第十九期-Minimum Cost Trees_3

往期

这一期主要写一个初代程序,能完整一些简单的例子

0 程序

import sys
import heapq
from collections import defaultdict

def read_graph_from_stdin():
    """ 从标准输入读取图数据并解析 """
    lines = sys.stdin.read().strip().split("\n")

    n = int(lines[0].strip())  # 节点数
    s = int(lines[1].strip())  # 源点
    k = int(lines[2].strip())  # 目标节点数
    terminals = list(map(int, lines[3].strip().split()))  # 目标节点列表
    D = int(lines[4].strip())  # 延迟约束
    m = int(lines[5].strip())  # 边数

    graph = defaultdict(list)

    for i in range(6, 6 + m):
        a, b, c, d = map(int, lines[i].strip().split())
        graph[a].append((b, c, d))  # 存储 a -> b 的边
        graph[b].append((a, c, d))  # 存储 b -> a 的边(无向图)

    return n, s, terminals, D, graph


def find_directed_tree(n, s, terminals, D, graph, avoid_edges=set()):
    """ 使用 Dijkstra/Prim 算法构造满足延迟约束的有向树,并尽量避免使用 avoid_edges """

    pq = [(0, s, -1)]  # (delay, node, parent)
    min_delay = {node: float('inf') for node in range(n)}
    min_delay[s] = 0
    tree_edges = []
    parent_map = {}

    while pq:
        delay, node, parent = heapq.heappop(pq)

        if parent != -1 and (parent, node) not in avoid_edges:
            tree_edges.append((parent, node))

        for neighbor, cost, edge_delay in sorted(graph[node], key=lambda x: x[2]):  # 按延迟排序
            if (node, neighbor) in avoid_edges:
                continue  # 避免使用第一棵树中的边

            new_delay = delay + edge_delay
            if new_delay <= D and new_delay < min_delay[neighbor]:
                min_delay[neighbor] = new_delay
                heapq.heappush(pq, (new_delay, neighbor, node))
                parent_map[neighbor] = node

    if not all(t in parent_map for t in terminals):
        return None  # 不能形成有效的树

    return tree_edges


def process_and_output():
    graph_data = read_graph_from_stdin()
    if not graph_data:
        return

    n, s, terminals, D, graph = graph_data

    # 找到第一棵树
    tree1_edges = find_directed_tree(n, s, terminals, D, graph)

    if not tree1_edges:
        print("无法找到符合要求的树")
        return

    used_edges = set(tree1_edges)

    # 尝试找到第二棵树(避免使用第一棵树的边)
    tree2_edges = find_directed_tree(n, s, terminals, D, graph, avoid_edges=used_edges)

    # 允许部分重用第一棵树的边
    if not tree2_edges:
        # print("无法找到完全不同的第二棵树,尝试部分重用第一棵树的边...")
        tree2_edges = find_directed_tree(n, s, terminals, D, graph, avoid_edges=set())

    if not tree2_edges:
        # print("无法找到第二棵符合要求的树")
        return

    print(2)  # 找到 2 棵树

    # 第一棵树
    print(len(tree1_edges))
    for a, b in tree1_edges:
        print(a, b)

    # 第二棵树
    print(len(tree2_edges))
    for a, b in tree2_edges:
        print(a, b)


if __name__ == "__main__":
    process_and_output()

在 PyCharm 终端运行

""" 从标准输入读取图数据并解析 """
lines = sys.stdin.read().strip().split("\n")

因为是标准数据输入,所以自己测试的时候,在终端中输入测试数据(输入完按 Ctrl+D 结束输入):
在这里插入图片描述

标准数据输入

  • 因为是标准数据输入,所以自己测试的时候,在终端中输入测试数据(输入完按 Ctrl+D 结束输入):
    因为是标准数据输入,所以自己测试的时候,在终端中输入测试数据(输入完按 Ctrl+D 结束输入):
    因为是标准数据输入,所以自己测试的时候,在终端中输入测试数据(输入完按 Ctrl+D 结束输入):
;