如何用Python bigtree将Polars/Pandas DataFrame转换为指定树结构？

核心思路 #

先对DataFrame按id分组聚合，收集每个节点对应的[val1, val2]列表并记录父节点关系，再用bigtree构建树结构，最后自定义打印逻辑输出目标格式。

Pandas 实现步骤 #

1. 导入依赖并准备数据

import pandas as pd
from bigtree import Node, build_tree_from_dict

# 替换为你的实际DataFrame
df = pd.DataFrame({
    "id": [1,1,1,2,2,3,3,3,4,4,4,4],
    "parent": [0,0,0,1,1,1,1,1,2,2,2,2],
    "val1": [10,20,30,40,50,60,70,80,90,100,110,120],
    "val2": [100,200,300,400,500,600,700,800,900,1000,1100,1200]
})

2. 分组聚合处理数据

按id分组，收集每个节点的[val1, val2]列表，同时提取父节点（同一id的父节点一致，取第一个值即可）：

grouped = df.groupby("id").agg(
    parent=("parent", "first"),
    vals=("val1", lambda x: list(zip(x, df.loc[x.index, "val2"])))
).reset_index()

3. 转换为bigtree兼容格式并构建树

将父节点为0的标记为根节点（parent_id=None），然后用build_tree_from_dict快速构建树：

# 转换为节点字典列表
node_dicts = grouped.to_dict("records")
for d in node_dicts:
    d["parent_id"] = d.pop("parent")
    if d["parent_id"] == 0:
        d["parent_id"] = None

# 构建树
root = build_tree_from_dict(node_dicts, node_id_attr="id", parent_id_attr="parent_id")

4. 自定义打印函数输出目标格式

def print_custom_tree(node, prefix="", is_last=True):
    # 格式化节点内容：id[[val1,val2], ...]
    vals_str = ", ".join([f"[{v1}, {v2}]" for v1, v2 in node.vals])
    node_str = f"{node.id}[{vals_str}]"
    
    # 打印当前节点
    if prefix:
        print(f"{prefix}{'|--' if not is_last else '`--'} {node_str}")
    else:
        print(node_str)
    
    # 递归处理子节点
    children = node.children
    for idx, child in enumerate(children):
        is_last_child = (idx == len(children) - 1)
        new_prefix = prefix + ("|  " if not is_last else "   ")
        print_custom_tree(child, new_prefix, is_last_child)

# 执行打印
print_custom_tree(root)

Polars 实现步骤 #

逻辑和Pandas一致，仅分组聚合部分略有差异：

1. 导入依赖并准备数据

import polars as pl
from bigtree import Node, build_tree_from_dict

# 替换为你的实际DataFrame
df = pl.DataFrame({
    "id": [1,1,1,2,2,3,3,3,4,4,4,4],
    "parent": [0,0,0,1,1,1,1,1,2,2,2,2],
    "val1": [10,20,30,40,50,60,70,80,90,100,110,120],
    "val2": [100,200,300,400,500,600,700,800,900,1000,1100,1200]
})

2. 分组聚合处理数据

grouped = df.group_by("id").agg(
    pl.col("parent").first().alias("parent"),
    pl.struct(["val1", "val2"]).list().alias("vals")
)

3. 转换为bigtree兼容格式并构建树

node_dicts = []
for row in grouped.iter_rows(named=True):
    # 提取val1和val2的元组列表
    vals_list = [(item["val1"], item["val2"]) for item in row["vals"]]
    node_dict = {
        "id": row["id"],
        "parent_id": row["parent"],
        "vals": vals_list
    }
    if node_dict["parent_id"] == 0:
        node_dict["parent_id"] = None
    node_dicts.append(node_dict)

# 构建树
root = build_tree_from_dict(node_dicts, node_id_attr="id", parent_id_attr="parent_id")

4. 调用相同的自定义打印函数

直接使用上面的print_custom_tree(root)即可输出目标格式。

输出效果 #

运行后会得到和预期一致的树结构：

1[[10, 100], [20, 200], [30, 300]]
|-- 2[[40, 400], [50, 500]]
|  |-- 4[[90, 900], [100, 1000], [110, 1100], [120, 1200]]
`-- 3[[60, 600], [70, 700], [80, 800]]

内容的提问来源于stack exchange，提问作者Galedon