過濾/篩選樹節點

邊城發表於2021-10-03

過濾/篩選樹節點

又是樹，是我跟樹槓上了嗎？—— 不，是樹的問題太多了！

? 相關文章推薦：

過濾和篩選是一個意思，都是 filter。

對於列表來說，過濾就是丟掉不需要的，留下需要的。但對於樹來說就得分情況了。

  • 如果想“過濾掉”（丟掉）某些節點，會把它的子節點一併拋棄，就像砍樹枝一樣，幹之不存，枝將焉附？這種情況多是去除不需要的子樹。

  • 如果是想“查詢”某些節點，會將找到的節點及其上溯到根的所有節點都保留下來。對於找到的節點，除了保留其完整路徑之外，對其子樹還有兩種處理方式：

    • 一種是“到此為止”，也就是說，如果其子樹中沒有符合條件的節點，那就不需要了，砍掉。需要定位到符合條件的節點以便後繼操作是採用這種方式，這也是最常用的查詢方式。
    • 另一種是保留其完整子樹。如果需要使用符合條件節點的子節點（比如選擇指定部門及其子部門）會採用這種方式。

過濾和查詢的主要區別在於：“過濾”通常會遇到不符合保留條件（或符合剔除條件）的節點就直接砍掉，不管其子樹中是否還存在符合保留條件的節點；而查詢則會一直找到葉節點上，只有整條路徑都沒有符合保留條件的節點，才會從其某個祖先節點上砍掉（祖先節點是否保留取決於其下是否存在符合保留條件的子孫節點）。

下面一樣一樣來。示例程式碼使用 TypeScript 編寫，示例資料來源於從列表生成樹 (JavaScript/TypeScript) 一文，同時使用該文中定義的節點型別介面：

interface TreeNode {
    id: number;
    parentId: number;
    label: string;
    children?: TreeNode[]
}

過濾掉不需要的節點

過濾掉不需要的節點，思路比較簡單：

  • 遍歷當前節點的所有子節點，需要的留，不需要的刪
  • 對留下的節點，通過遞迴進行過濾

按此思路，TypeScript 程式碼是

/**
 * @param nodes 要過濾的樹節點集（多根）
 * @param predicate 過濾條件，返回 `true` 保留
 * @returns 過濾後的樹節點集
 */
function filterTree(
    nodes: TreeNode[] | undefined,
    predicate: (node: TreeNode) => boolean
): TreeNode[] | undefined {
    if (!nodes?.length) { return nodes; }

    // 直接使用 Array 的 filter 可以過濾當層節點
    return nodes.filter(it => {
        // 不符合條件的直接砍掉
        if (!predicate(it)) {
            return false;
        }

        // 符合條件的保留，並且需要遞迴處理其子節點
        it.children = filterTree(it.children, predicate);
        return true;
    });
}

如果對示例資料（見前文）進行過濾，僅保留 id 是偶數的節點，那結果是

flowchart LR
%%{ init: { "theme": "forest" } }%%

S(("Virtual\nRoot"))
S --> N6
S --> N10

N6("6 | P6mtcgfCD")
N6 --> N8("8 | m6o5UsytQ0")
N10("10 | lhDGTNeeSxLNJ")
N6 --> N14("14 | ysYwG8EFLAu1a")
N10 --> N16("16 | RKuQs4ki65wo")

不過這個 filterTree 有點小瑕疵：

  1. 遞迴呼叫時還需要傳入 predicate，有點繁瑣
  2. 傳入引數應該限制在 TreeNode[] 型別上，新增 undefined 只是為了簡化遞迴呼叫（不用先判空）

處理起來也簡單，加一層介面封裝一下（門面模式）：

/**
 * @param nodes 要過濾的樹節點集（多根）
 * @param predicate 過濾條件，返回 `true` 保留
 * @returns 過濾後的樹節點集
 */
function filterTree(
    nodes: TreeNode[],
    predicate: (node: TreeNode) => boolean
): TreeNode[] {
    return filter(nodes) ?? [];

    // 原 filterTree，更名，並刪除 predicate 引數
    function filter(nodes: TreeNode[] | undefined): TreeNode[] | undefined {
        if (!nodes?.length) { return nodes; }

        return nodes.filter(it => {
            if (!predicate(it)) {
                return false;
            }
            // 遞迴呼叫不需要再傳入 predicate
            it.children = filter(it.children);
            return true;
        });
    }
}

實際使用的時候，可能傳入的可能是單根樹 (TreeNode)，也有可能是多根 (TreeNode[])，那可以寫個過載：

function filterTree(node: TreeNode, predicate: (node: TreeNode) => boolean): TreeNode;
function filterTree(nodes: TreeNode[], predicate: (node: TreeNode) => boolean): TreeNode[];
function filterTree(
    tree: TreeNode | TreeNode[],
    predicate: (node: TreeNode) => boolean
): TreeNode | TreeNode[] {
    if (Array.isArray(tree)) {
        return filter(tree) ?? [];
    } else {
        tree.children = filter(tree.children);
        return tree;
    }

    function filter(...) { ... }
}

查詢節點（不含完整子樹）

查詢節點就要稍微複雜了點了，因為需要保留路徑。判斷當前節點是否可以刪除需要對自己情況進行判斷之外，還取決於其所有子孫節點是否可以刪除。與前面“過濾掉”的邏輯相比，有兩點變化：

  1. 不管當前節點是否保留，均需要遞迴向下，把子孫中符合條件的節點都找出來
  2. 只要子孫中存在符合條件的節點，當前節點就應該保留。

這樣處理後的節點，所有葉節點都應該符合查詢條件。比如在示例資料中按 id 參整除 6 來查詢節點，結果是：

flowchart LR
%%{ init: { "theme": "forest" } }%%
classDef found fill:#ffeeee,stroke:#cc6666;

S(("Virtual\nRoot")) --> N1
S --> N6:::found;

N1("1 | 8WUg35y")
N1 --> N4("4 | IYkxXlhmU12x")
N4 --> N5("5 | p2Luabg9mK2")
N6("6 | P6mtcgfCD")
N1 --> N7("7 | yluJgpnqKthR")
N7 --> N12("12 | 5W6vy0EuvOjN"):::found
N5 --> N13("13 | LbpWq")
N13 --> N18("18 | 03X6e4UT"):::found

根據上面的邏輯，寫一個 findTreeNode()：

function findTreeNode(node: TreeNode, predicate: (node: TreeNode) => boolean): TreeNode;
function findTreeNode(nodes: TreeNode[], predicate: (node: TreeNode) => boolean): TreeNode[];
function findTreeNode(
    tree: TreeNode | TreeNode[],
    predicate: (node: TreeNode) => boolean
): TreeNode | TreeNode[] {
    if (Array.isArray(tree)) {
        return filter(tree) ?? [];
    } else {
        tree.children = filter(tree.children);
        return tree;
    }

    function filter(nodes: TreeNode[] | undefined): TreeNode[] | undefined {
        if (!nodes?.length) { return nodes; }
        return nodes.filter(it => {
            // 先篩選子樹，如果子樹中沒有符合條件的，children 會是 [] 或 undefined
            const children = filter(it.children);
            // 根據當前節點情況和子樹篩選結果判斷是否保留當前節點
            if (predicate(it) || children?.length) {
                // 如果要保留，children 應該用篩選出來的那個；不保留的話就不 care 子節點了
                it.children = children;
                return true;
            }
            return false;
        });
    }
}

下面把程式碼修改下，在結果中保留子樹。

查詢節點（含完整子樹）

這個思路跟最上面那個“剔除”的思路正好相反，

  • 遇到符合條件的節點，直接保留整棵子樹，也不需要遞迴去處理了
  • 不符合條件的節點，遞迴進去繼續找

既然都是查詢，可以給 findTreeNode() 新增一個 keepSubTree: boolean 引數來擴充套件函式功能。介面部分改變如下：

function findTreeNode(
    node: TreeNode,
    predicate: (node: TreeNode) => boolean,
    keepSubTree?: boolean  // <--
): TreeNode;
function findTreeNode(
    nodes: TreeNode[],
    predicate: (node: TreeNode) => boolean,
    keepSubTree?: boolean  // <--
): TreeNode[];
function findTreeNode(
    tree: TreeNode | TreeNode[],
    predicate: (node: TreeNode) => boolean,
    keepSubTree: boolean = false  // <--
): TreeNode | TreeNode[] {
    ...
}

然後需要修改的地方主要是 Array.prototype.filter 回撥函式，可以先把原來的箭頭函式提取出來，命名為 filterWithoutSubTree()。

提取函式（動畫）

然後再寫一個 filterWithSubTree() 處理函式。根據 keepSubTree 的值來決定使用哪一個過濾器。關鍵程式碼如下：

function findTreeNode(...): TreeNode | TreeNode[] {
    const filterHandler = keepSubTree ? filterWithSubTree : filterWithoutSubTree;
    //    ^^^^^^^^^^^^^

    if (Array.isArray(tree)) { ... } else { ... }

    function filter(nodes: TreeNode[] | undefined): TreeNode[] | undefined {
        if (!nodes?.length) { return nodes; }
        return nodes.filter(filterHandler);
        //                  ^^^^^^^^^^^^^
    }

    function filterWithSubTree(it: TreeNode): boolean {
        // 如果符合條件，保留整棵子樹，不需要遞迴進去
        if (predicate(it)) { return true; }

        // 否則根據子孫節點的情況來決定是否需要保留當前節點（作為路徑節點）
        it.children = filter(it.children);
        return !!it.children?.length;
    }

    function filterWithoutSubTree(it: TreeNode): boolean {
        ...
    }
}

含完整子樹的查詢結果示例（查詢條件：it => it.id % 4 === 0）如下圖：

flowchart LR
%%{ init: { "theme": "forest" } }%%
classDef found fill:#ffeeee,stroke:#cc6666;
classDef subs fill:#ffffff;

S(("Virtual\nRoot")) --> N1
S --> N6
S --> N10

N1("1 | 8WUg35y")
N1 --> N4("4 | IYkxXlhmU12x"):::found
N4 --> N5("5 | p2Luabg9mK2"):::subs
N6("6 | P6mtcgfCD")
N1 --> N7("7 | yluJgpnqKthR")
N6 --> N8("8 | m6o5UsytQ0"):::found
N10("10 | lhDGTNeeSxLNJ")
N7 --> N12("12 | 5W6vy0EuvOjN"):::found
N5 --> N13("13 | LbpWq"):::subs
N10 --> N16("16 | RKuQs4ki65wo"):::found
N13 --> N18("18 | 03X6e4UT"):::subs
N7 --> N19("19 | LTJTeF")
N19 --> N20("20 | 3rqUqE3MLShh"):::found

相關文章