如何通过PostgreSQL递归查询实现二叉树节点按规则着色？

问题场景 #

我们在PostgreSQL中用一张表存储二叉树结构，每行对应一个节点，包含node（节点标识）、parent（父节点）、color（颜色）字段。其中叶子节点的color字段非空，取值为green或red，需要按照以下规则为所有非叶子节点计算颜色：

• 若父节点仅有一个子节点，父节点颜色与子节点相同；
• 若父节点有两个子节点且颜色一致，父节点颜色与子节点相同；
• 若父节点有两个子节点且颜色不同，父节点颜色为gray。

解决方案 #

使用PostgreSQL的递归CTE（Common Table Expression）可以实现从叶子节点向上回溯，逐层计算父节点的颜色。递归CTE分为基例（已知颜色的叶子节点）和递归步骤（向上推导父节点颜色）两部分。

完整查询代码 #

WITH RECURSIVE tree_color AS (
    -- 基例：叶子节点，直接使用已有颜色
    SELECT node, parent, color
    FROM family
    WHERE color IS NOT NULL

    UNION ALL

    -- 递归步骤：基于子节点颜色计算父节点颜色
    SELECT 
        f.node,
        f.parent,
        CASE
            -- 子节点数量为1，继承子节点颜色
            WHEN COUNT(tc.node) = 1 THEN MAX(tc.color)
            -- 两个子节点颜色一致，继承该颜色
            WHEN COUNT(DISTINCT tc.color) = 1 THEN MAX(tc.color)
            -- 两个子节点颜色不同，设为gray
            ELSE 'gray'
        END AS color
    FROM family f
    INNER JOIN tree_color tc ON f.node = tc.parent
    WHERE f.color IS NULL
    GROUP BY f.node, f.parent
)
-- 输出所有节点的最终着色结果
SELECT node, parent, color
FROM tree_color
ORDER BY node;

代码解释 #

1. 基例部分：首先筛选出所有叶子节点（color IS NOT NULL），这是递归计算的起始点，因为它们的颜色是已知的固定值。
2. 递归步骤：
   • 通过INNER JOIN将未着色的父节点（f.color IS NULL）与已经计算出颜色的子节点（tree_color）关联。
   • 对每个父节点进行分组，统计其子节点的数量和颜色的唯一值数量：
     • 若子节点数量为1，直接取该子节点的颜色；
     • 若两个子节点的颜色唯一值数量为1（即颜色相同），取该颜色；
     • 若两个子节点颜色不同（唯一值数量为2），则将父节点颜色设为gray。
3. 最终输出：从递归结果集中取出所有节点的信息，按节点名称排序后输出，得到整棵树的完整着色结果。

测试结果 #

使用题目提供的临时表数据运行上述查询，会得到如下符合规则的着色结果：

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