LC 6223.移除子樹後的二叉樹高度(dp ||dfs序)

思路1.dfs序

顯然子樹問題可以用dfs序轉成區間問題。

本題可以轉的原因，删除一個子樹，相當删除一個子樹，而除了該子樹外的其他點都可以從根節點走到(題目保證根不會被删)。是以每次詢問就剩下一個字首和字尾，取他們兩的最大值就是答案，這個值定義為每個結點的深度。

時間複雜度：

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    int n = 0, clk = 0;
    vector<int> A, L, R;

    void dfs(TreeNode *node, int d) {
        // 沒有輸入 n，隻好動态計算 n 的大小并擴充清單...
        int idx = node->val;
        n = max(idx, n);
        while (A.size() <= n) A.push_back(0), L.push_back(0), R.push_back(0);

        // node 是第 clk 個被通路的點
        clk++;
        // A[i] 表示第 i 個被通路的點的深度
        A[clk] = d;
        // L[i] 表示第 i 個點的子樹對應的連續區間的左端點
        L[idx] = clk;

        // DFS 子樹
        if (node->left != nullptr) dfs(node->left, d + 1);
        if (node->right != nullptr) dfs(node->right, d + 1);

        // R[i] 表示第 i 個點的子樹對應的連續區間的右端點
        R[idx] = clk;
    }

public:
    vector<int> treeQueries(TreeNode* root, vector<int>& queries) {
        dfs(root, 0);

        // f[i] 表示 max(A[1], A[2], ..., A[i])
        // g[i] 表示 max(A[n], A[n - 1], ..., A[i])
        vector<int> f(n + 2), g(n + 2);
        for (int i = 1; i <= n; i++) f[i] = max(f[i - 1], A[i]);
        for (int i = n; i > 0; i--) g[i] = max(g[i + 1], A[i]);

        vector<int> ans;
        // 樹上詢問轉為區間詢問處理
        for (int x : queries) ans.push_back(max(f[L[x] - 1], g[R[x] + 1]));
        return ans;
    }
};      

思路2 dp

先預處理出每個子樹的最大高度，這個可以遞歸。

class Solution {
public:
    vector<int> treeQueries(TreeNode *root, vector<int> &queries) {
        unordered_map<TreeNode*, int> height; // 每棵子樹的高度
        function<int(TreeNode*)> get_height = [&](TreeNode *node) -> int {
            return node ? height[node] = 1 + max(get_height(node->left), get_height(node->right)) : 0;
        };
        get_height(root);

        int res[height.size() + 1]; // 每個節點的答案
        function<void(TreeNode*, int, int)> dfs = [&](TreeNode *node, int depth, int rest_h) {
            if (node == nullptr) return;
            ++depth;
            res[node->val] = rest_h;
            dfs(node->left, depth, max(rest_h, depth + height[node->right]));
            dfs(node->right, depth, max(rest_h, depth + height[node->left]));
        };
        dfs(root, -1, 0);

        for (auto &q : queries) q = res[q];
        return queries;
    }
};