tree_dp.cpp

/*
Tree Dynamic Programming
Mathematical Foundation: dp[u] = f(dp[v1], dp[v2], ...) for all children v of u
Rerooting: Compute answer for all nodes as root in O(n)
Time: O(n), Space: O(n)
*/

#include <bits/stdc++.h>
using namespace std;

struct TreeNode {
    int val;
    TreeNode *left, *right;
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};

// Maximum Path Sum
// LeetCode: 124. Binary Tree Maximum Path Sum
// https://leetcode.com/problems/binary-tree-maximum-path-sum/
// Related:
// 543. Diameter of Binary Tree
// https://leetcode.com/problems/diameter-of-binary-tree/
// 687. Longest Univalue Path
// https://leetcode.com/problems/longest-univalue-path/
// 1245. Tree Diameter
// https://leetcode.com/problems/tree-diameter/
// 1372. Longest ZigZag Path in a Binary Tree
// https://leetcode.com/problems/longest-zigzag-path-in-a-binary-tree/
// 2246. Longest Path With Different Adjacent Characters
// https://leetcode.com/problems/longest-path-with-different-adjacent-characters/
int maxPathSum(TreeNode* root) {
    int res = INT_MIN;
    function<int(TreeNode*)> dfs = [&](TreeNode* node) {
        if (!node) return 0;
        int l = max(0, dfs(node->left));
        int r = max(0, dfs(node->right));
        res = max(res, node->val + l + r);
        return node->val + max(l, r);
    };
    dfs(root);
    return res;
}

// House Robber III
// LeetCode: 337. House Robber III
// https://leetcode.com/problems/house-robber-iii/
// Related:
// 198. House Robber
// https://leetcode.com/problems/house-robber/
// 213. House Robber II
// https://leetcode.com/problems/house-robber-ii/
// 740. Delete and Earn
// https://leetcode.com/problems/delete-and-earn/
// 124. Binary Tree Maximum Path Sum
// https://leetcode.com/problems/binary-tree-maximum-path-sum/
// 543. Diameter of Binary Tree
// https://leetcode.com/problems/diameter-of-binary-tree/
int rob(TreeNode* root) {
    function<pair<int,int>(TreeNode*)> dfs = [&](TreeNode* node) -> pair<int,int> {
        if (!node) return {0, 0};
        auto [ll, lr] = dfs(node->left);
        auto [rl, rr] = dfs(node->right);
        return {node->val + lr + rr, max(ll, lr) + max(rl, rr)};
    };
    auto [rob, skip] = dfs(root);
    return max(rob, skip);
}

// Tree Diameter (General Tree)
// LeetCode: 1245. Tree Diameter
// https://leetcode.com/problems/tree-diameter/
// Related:
// 543. Diameter of Binary Tree
// https://leetcode.com/problems/diameter-of-binary-tree/
// 124. Binary Tree Maximum Path Sum
// https://leetcode.com/problems/binary-tree-maximum-path-sum/
// 687. Longest Univalue Path
// https://leetcode.com/problems/longest-univalue-path/
// 2246. Longest Path With Different Adjacent Characters
// https://leetcode.com/problems/longest-path-with-different-adjacent-characters/
// 310. Minimum Height Trees
// https://leetcode.com/problems/minimum-height-trees/
int treeDiameter(vector<vector<int>>& adj) {
    int n = adj.size(), res = 0;
    function<int(int, int)> dfs = [&](int u, int p) {
        int mx1 = 0, mx2 = 0;
        for (int v : adj[u]) if (v != p) {
            int d = dfs(v, u);
            if (d > mx1) mx2 = mx1, mx1 = d;
            else if (d > mx2) mx2 = d;
        }
        res = max(res, mx1 + mx2);
        return mx1 + 1;
    };
    dfs(0, -1);
    return res;
}

// Rerooting DP - Sum of Distances
// LeetCode: 834. Sum of Distances in Tree
// https://leetcode.com/problems/sum-of-distances-in-tree/
// Related:
// 1245. Tree Diameter
// https://leetcode.com/problems/tree-diameter/
// 310. Minimum Height Trees
// https://leetcode.com/problems/minimum-height-trees/
// 2581. Count Number of Possible Root Nodes
// https://leetcode.com/problems/count-number-of-possible-root-nodes/
// 2920. Maximum Points After Collecting Coins From All Nodes
// https://leetcode.com/problems/maximum-points-after-collecting-coins-from-all-nodes/
// 968. Binary Tree Cameras
// https://leetcode.com/problems/binary-tree-cameras/
vector<long long> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
    vector<vector<int>> adj(n);
    for (auto& e : edges) {
        adj[e[0]].push_back(e[1]);
        adj[e[1]].push_back(e[0]);
    }
    
    vector<int> cnt(n);
    vector<long long> res(n);
    
    function<void(int, int)> dfs1 = [&](int u, int p) {
        cnt[u] = 1;
        for (int v : adj[u]) if (v != p) {
            dfs1(v, u);
            cnt[u] += cnt[v];
            res[u] += res[v] + cnt[v];
        }
    };
    
    function<void(int, int)> dfs2 = [&](int u, int p) {
        for (int v : adj[u]) if (v != p) {
            res[v] = res[u] - cnt[v] + (n - cnt[v]);
            dfs2(v, u);
        }
    };
    
    dfs1(0, -1);
    dfs2(0, -1);
    return res;
}

// Count Good Nodes in Binary Tree
// LeetCode: 1448. Count Good Nodes in Binary Tree
// https://leetcode.com/problems/count-good-nodes-in-binary-tree/
// Related:
// 104. Maximum Depth of Binary Tree
// https://leetcode.com/problems/maximum-depth-of-binary-tree/
// 1026. Maximum Difference Between Node and Ancestor
// https://leetcode.com/problems/maximum-difference-between-node-and-ancestor/
// 98. Validate Binary Search Tree
// https://leetcode.com/problems/validate-binary-search-tree/
// 250. Count Univalue Subtrees
// https://leetcode.com/problems/count-univalue-subtrees/
// 124. Binary Tree Maximum Path Sum
// https://leetcode.com/problems/binary-tree-maximum-path-sum/
int goodNodes(TreeNode* root) {
    function<int(TreeNode*, int)> dfs = [&](TreeNode* node, int mx) {
        if (!node) return 0;
        int res = node->val >= mx ? 1 : 0;
        mx = max(mx, node->val);
        return res + dfs(node->left, mx) + dfs(node->right, mx);
    };
    return dfs(root, INT_MIN);
}

// All Possible Full Binary Trees
// LeetCode: 894. All Possible Full Binary Trees
// https://leetcode.com/problems/all-possible-full-binary-trees/
// Related:
// 95. Unique Binary Search Trees II
// https://leetcode.com/problems/unique-binary-search-trees-ii/
// 96. Unique Binary Search Trees
// https://leetcode.com/problems/unique-binary-search-trees/
// 1214. Two Sum BSTs
// https://leetcode.com/problems/two-sum-bsts/
// 1379. Find a Corresponding Node of a Binary Tree in a Clone of That Tree
// https://leetcode.com/problems/find-a-corresponding-node-of-a-binary-tree-in-a-clone-of-that-tree/
// 337. House Robber III
// https://leetcode.com/problems/house-robber-iii/
vector<TreeNode*> allPossibleFBT(int n) {
    if (n % 2 == 0) return {};
    if (n == 1) return {new TreeNode(0)};
    vector<TreeNode*> res;
    for (int i = 1; i < n; i += 2) {
        auto left = allPossibleFBT(i);
        auto right = allPossibleFBT(n - 1 - i);
        for (auto l : left) for (auto r : right) {
            TreeNode* root = new TreeNode(0);
            root->left = l;
            root->right = r;
            res.push_back(root);
        }
    }
    return res;
}