Sympy如何提取表达式树指定层级的Sympy对象列表

我来帮你搞定在SymPy中提取表达式树任意层级节点的需求——不用依赖srepr()的字符串，直接拿到SymPy对象列表。

基础：获取第一层节点 #

你已经知道用Add.make_args(expr)，其实SymPy的所有表达式对象都自带args属性，直接用它就能拿到当前节点的直接子节点（也就是第一层元素）：

from sympy import *
x,y,z = symbols('x y z')
expr = sin(x*y)/2 - x**2 + 1/y

# 获取第一层节点
first_level = list(expr.args)
print(first_level)
# 输出：[-x**2, sin(x*y)/2, 1/y]

注意：SymPy的Add会自动对项排序，所以顺序可能和你手写的原式略有不同，如果要严格保持输入顺序，可以用UnevaluatedExpr包裹每个项再组合。

深入：获取任意深层级的节点 #

要拿到更深层级的节点，本质就是递归遍历表达式树的每个节点，或者通过路径索引定位到目标节点后再取其子节点。

方法1：按路径索引特定节点的子节点

如果你明确知道要找的节点在树中的位置，可以通过链式调用args来定位：
比如你要获取first_level[1]（也就是sin(x*y)/2）的子节点：

# 定位到sin(x*y)/2，再取它的子节点
second_level_0 = list(expr.args[1].args)
print(second_level_0)
# 输出：[sin(x*y), 1/2]

如果你的目标是得到[sin(x*y), 2]，只需要对第二个元素取倒数即可：

second_level_0_adjusted = [second_level_0[0], second_level_0[1]**-1]
print(second_level_0_adjusted)
# 输出：[sin(x*y), 2]

方法2：批量获取指定层级的所有节点

如果需要一次性拿到某一层级的所有节点，可以写一个递归函数来遍历整个表达式树：

def get_all_nodes_at_level(expr, target_level, current_level=1):
    """
    获取表达式树中指定层级的所有节点
    - target_level: 要获取的层级（第一层为1）
    - current_level: 当前递归到的层级，默认不用传参
    """
    if current_level == target_level:
        return [expr]
    nodes = []
    # 遍历当前节点的所有子节点，递归收集下一层的节点
    for arg in expr.args:
        nodes.extend(get_all_nodes_at_level(arg, target_level, current_level + 1))
    return nodes

# 获取第二层所有节点
second_level = get_all_nodes_at_level(expr, 2)
print(second_level)
# 输出：[x**2, -1, sin(x*y), 1/2, y, -1]

# 获取第三层所有节点
third_level = get_all_nodes_at_level(expr, 3)
print(third_level)
# 输出：[x, 2, x*y]

核心逻辑总结 #

SymPy的表达式树遵循操作符作为节点，操作数作为子节点的结构：

• 叶子节点（比如x、2、Rational(1,2)）的args是空元组
• 非叶子节点（比如Add、Mul、sin）的args就是它的直接子节点列表
• 不管层级多深，只要通过args链式访问或递归遍历，就能拿到任意位置的节点对象

内容的提问来源于stack exchange，提问作者Uwe.Schneider