用Julia或Python绘制公平硬币抛掷3次结果的图及概率计算

实现思路 #

动态生成硬币抛掷的状态树邻接矩阵，适配任意抛掷次数，自动计算每个节点的概率。状态树以空字符串为根节点（0次抛掷），每一层对应k次抛掷的所有序列结果，非叶子节点均有H、T两个子节点。

完整代码 #

using GraphRecipes, Plots
gr()
default(size=(1000, 600))

function coin_toss_tree(n::Int)
    total_nodes = 2^(n+1) - 1
    adj_matrix = zeros(Int, total_nodes, total_nodes)
    node_names = String[]
    probabilities = Float64[]

    for k in 0:n
        layer_size = 2^k
        start_idx = iszero(k) ? 1 : sum(2^i for i in 0:k-1) + 1
        end_idx = start_idx + layer_size - 1

        for i in start_idx:end_idx
            if k == 0
                name = ""
                prob = 1.0
            else
                parent_idx = (i - start_idx) ÷ 2 + sum(2^j for j in 0:k-2) + 1
                parent_name = node_names[parent_idx]
                name = parent_name * ((i - start_idx) % 2 == 0 ? "H" : "T")
                prob = 1.0 / 2^k
            end
            push!(node_names, name)
            push!(probabilities, prob)

            if k < n
                left_child = end_idx + (i - start_idx)*2 + 1
                right_child = left_child + 1
                adj_matrix[i, left_child] = 1
                adj_matrix[i, right_child] = 1
            end
        end
    end

    p = graphplot(adj_matrix,
                  layout=:tree,
                  fontsize=10,
                  names=node_names,
                  nodeshape=:circle,
                  nodecolor=:lightblue,
                  nodesize=0.15,
                  edgecolor=:gray,
                  root=1,
                  branchlength=0.5)
    
    annotate!([(p.plots[1].series_list[1].x[i], p.plots[1].series_list[1].y[i] + 0.05, 
                text(round(probabilities[i], digits=4), 8, :black)) for i in 1:total_nodes])
    
    title!("公平硬币抛掷 $n 次状态树（含概率）")
    return p, probabilities
end

# 运行示例：3次抛掷
plt, probs = coin_toss_tree(3)
display(plt)

关键优化点 #

• 动态生成邻接矩阵，无需手动编写，支持扩展到任意抛掷次数
• 使用:tree布局让图形呈现清晰的层级结构，避免原代码中节点混乱的问题
• 自动计算每个节点的概率：k次抛掷的节点概率为1/2^k
• 添加概率标注，直观展示各状态的发生概率

实现思路 #

用networkx构建有向状态树，matplotlib配合graphviz绘制树状图，自动生成节点、边及概率属性，支持任意抛掷次数扩展。

完整代码 #

import networkx as nx
import matplotlib.pyplot as plt

def coin_toss_tree(n):
    G = nx.DiGraph()
    # 根节点：0次抛掷状态
    root_node = ""
    G.add_node(root_node, prob=1.0)
    current_layer_nodes = [root_node]

    for k in range(1, n+1):
        next_layer_nodes = []
        for parent_node in current_layer_nodes:
            # 生成H、T子节点
            h_child = parent_node + "H"
            t_child = parent_node + "T"
            node_prob = 1.0 / (2 ** k)
            
            G.add_node(h_child, prob=node_prob)
            G.add_node(t_child, prob=node_prob)
            G.add_edge(parent_node, h_child)
            G.add_edge(parent_node, t_child)
            
            next_layer_nodes.extend([h_child, t_child])
        current_layer_nodes = next_layer_nodes

    # 使用dot布局生成树状结构
    pos = nx.nx_agraph.graphviz_layout(G, prog="dot")
    plt.figure(figsize=(12, 8))

    # 绘制节点、边
    nx.draw_networkx_nodes(G, pos, node_size=2000, node_color="lightblue")
    nx.draw_networkx_edges(G, pos, edge_color="gray", arrows=True, arrowsize=15)

    # 添加节点名称和概率标签
    node_labels = {node: node for node in G.nodes()}
    nx.draw_networkx_labels(G, pos, node_labels, font_size=12)
    
    prob_labels = {node: f"{G.nodes[node]['prob']:.4f}" for node in G.nodes()}
    nx.draw_networkx_labels(G, pos, prob_labels, font_size=8, label_pos=0.3)

    plt.title(f"公平硬币抛掷 {n} 次状态树（含概率）")
    plt.axis("off")
    plt.show()

# 运行示例：3次抛掷
coin_toss_tree(3)

依赖安装 #

执行代码前需安装依赖包：

pip install networkx matplotlib graphviz

注意：系统需提前安装graphviz软件（Ubuntu用sudo apt install graphviz，Windows/macOS可从官网下载安装）。

内容的提问来源于stack exchange，提问作者Freya the Goddess