Mastering Ancestor Tracking in Trees

Tree problems are crucial for reaching the Above Expert level on Codeforces and the Guardian badge on LeetCode. Some fundamental tree topics include:

*DFS

*BFS

*DP on Trees (In/Out DP)

*Rerooting

*Binary Lifting

*Ancestor Tracking

In this blog, we will focus on Ancestor Tracking in trees.

Understanding the Problem Statement

You are given a tree with constraints that must be maintained on a downward path.

Consider the following problem from LeetCode:

Longest Special Path

Problem Statement:

You are given an undirected tree rooted at node 0 with n nodes numbered from 0 to n — 1. The tree is represented by a 2D array edges of length n — 1, where:

edges[i] = [ui, vi, lengthi] represents an edge between nodes ui and vi with length lengthi.

An integer array nums is given, where nums[i] represents the value at node i.

A special path is a downward path from an ancestor node to a descendant node such that all values in the path are unique.

Return: An array result of size 2, where:

result[0] is the length of the longest special path.

result[1] is the minimum number of nodes in all possible longest special paths.

Key Observation

Since we must maintain a constraint (unique node values in a downward path), we can directly apply the ancestor property.

Template for Solving Ancestor Tracking Problems

Follow these four steps:

1.Choose a data structure (e.g., map, multiset) to store ancestor properties.

2.Insert the current node into the data structure.

3.Compute the answer from the data structure while ensuring the given constraint is not violated.

4.Backtrack by removing the current node from the data structure.

C++ Solution

class Solution {
public:
    // Choosing the data structure
    multiset<int> st;
    map<int, int> mp;
    map<int, vector<int>> lastseen;
    vector<int> ans;
    
    void dfs(int u, int p, vector<array<int, 2>> adj[], vector<int>& nums, vector<int>& depth, vector<int>& length){
        for(auto &[v, w] : adj[u]){
            if(v == p) continue;
            depth[v] = depth[u] + 1;
            length[v] = length[u] + w;
            dfs(v, u, adj, nums, depth, length);
        }
    }
    
    void solve(int u, int p, vector<array<int, 2>> adj[], vector<int>& nums, vector<int>& depth, vector<int>& length, int lvl){
        // Insert the node into the data structure
        if(!lastseen[nums[u]].empty()){
            st.insert(depth[lastseen[nums[u]].back()]);
        }
        lastseen[nums[u]].push_back(u);
        mp[lvl] = u;

        // Finding the answer for the current node
        int safeNode;
        if(st.empty()){
            safeNode = mp[0];
        } else {
            int conflictDepth = *st.rbegin();
            int safeLvl = conflictDepth + 1;
            safeNode = mp[safeLvl];
        }
        
        int currLength = length[u] - length[safeNode];
        int currCount = depth[u] - depth[safeNode] + 1;
        
        if(currLength > ans[0]){
            ans[0] = currLength;
            ans[1] = currCount;
        } else if(currLength == ans[0]){
            ans[1] = min(ans[1], currCount);
        }
        
        for(auto &[v, w] : adj[u]){
            if(v == p) continue;
            solve(v, u, adj, nums, depth, length, lvl + 1);
        }

        // Backtracking
        mp.erase(lvl);
        lastseen[nums[u]].pop_back();
        if(!lastseen[nums[u]].empty()){
            st.erase(st.find(depth[lastseen[nums[u]].back()]));
        }
    }
    
    vector<int> longestSpecialPath(vector<vector<int>>& edges, vector<int>& nums) {
        int n = nums.size();
        vector<int> depth(n, 0), length(n, 0);
        ans = {0, INT_MAX};
        vector<array<int, 2>> adj[n];
        
        for(int i = 0; i < edges.size(); i++){
            int u = edges[i][0], v = edges[i][1], w = edges[i][2];
            adj[u].push_back({v, w});
            adj[v].push_back({u, w});
        }
        
        dfs(0, -1, adj, nums, depth, length);
        solve(0, -1, adj, nums, depth, length, 0);
        return ans;
    }
};

Further Exploration

If you have understood the problem, try solving Longest Special Path II. It requires only a minor modification to the above approach!

If you're stuck, drop a comment, and I'll be happy to provide hints! Please also provide some more problems on the same idea in the comment section :)

Happy Coding!