Construct ancestor matrix from a binary tree

Problem

Given a binary tree where all node values are from 0 to n-1, construct an ancestor matrix mat[n][n] such that:

mat[i][j] = 1 if i is an ancestor of j.
mat[i][j] = 0 otherwise.

An ancestor of a node is defined as any node that is found along the path from the root to that node in the binary tree.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


 Input:
 Root =
        0
       / \
      1   2
 Output:
 Ancestor Matrix:
 [0 1 1]
 [0 0 0]
 [0 0 0]
 Explanation:
 - Node 0 is the ancestor of both 1 and 2, so `mat[0][1] = 1` and `mat[0][2] = 1`.
 - Nodes 1 and 2 are leaf nodes and do not have descendants, so their rows are all zeros.

Example 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17


 Input:
 Root =
        0
       / \
      1   2
         / \
        3   4
 Output:
 Ancestor Matrix:
 [0 1 1 1 1]
 [0 0 0 0 0]
 [0 0 0 1 1]
 [0 0 0 0 0]
 [0 0 0 0 0]
 Explanation:
 - Node 0 is the ancestor of nodes 1, 2, 3, and 4; hence `mat[0][1] = mat[0][2] = mat[0][3] = mat[0][4] = 1`.
 - Node 2 is the ancestor of nodes 3 and 4; hence `mat[2][3] = mat[2][4] = 1`.

Solution

Method 1 - Using DFS

To construct the ancestor matrix, traverse the binary tree and track ancestors of each node:

Use a depth-first traversal of the tree.
Maintain a running list (or set) of current ancestors for the node being visited.
For every node, mark entries in the matrix to indicate that all its ancestors are indeed ancestors of the node.

Approach

Initialisation:
- Create an n x n matrix filled with zeros.
- Traverse the binary tree starting from its root.
DFS Traversal:
- For each node, maintain a list of ancestors passed down recursively to children.
- At each node, populate the matrix by marking 1 for all positions [ancestor][currentNode] in the matrix.
Helper Function for DFS:
- Input: Current node and list of ancestors.
- Mark respective entries in the matrix.
- Append the current node to the ancestor list and recursively process left and right children.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45


public class Solution {
    public int[][] constructAncestorMatrix(TreeNode root, int n) {
        int[][] matrix = new int[n][n];
        List<Integer> ancestors = new ArrayList<>();
        dfs(root, matrix, ancestors);
        return matrix;
    }

    private void dfs(TreeNode node, int[][] matrix, List<Integer> ancestors) {
        if (node == null) return;

        // Mark current node's ancestors in the matrix
        for (int ancestor : ancestors) {
            matrix[ancestor][node.val] = 1;
        }

        // Add current node to the ancestors list and recurse
        ancestors.add(node.val);
        dfs(node.left, matrix, ancestors);
        dfs(node.right, matrix, ancestors);

        // Remove current node from ancestors list (backtracking)
        ancestors.remove(ancestors.size() - 1);
    }

    public static void main(String[] args) {
        Solution solution = new Solution();
        
        TreeNode root = new TreeNode(0);
        root.left = new TreeNode(1);
        root.right = new TreeNode(2);
        root.right.left = new TreeNode(3);
        root.right.right = new TreeNode(4);

        int n = 5;
        int[][] ancestorMatrix = solution.constructAncestorMatrix(root, n);

        for (int[] row : ancestorMatrix) {
            for (int val : row) {
                System.out.print(val + " ");
            }
            System.out.println();
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37


class Solution:
    def constructAncestorMatrix(self, root, n):
        # Initialise matrix and an empty ancestor list
        matrix = [[0 for _ in range(n)] for _ in range(n)]
        self.dfs(root, matrix, [])
        return matrix

    def dfs(self, node, matrix, ancestors):
        if not node:
            return

        # Mark current node's ancestors in the matrix
        for ancestor in ancestors:
            matrix[ancestor][node.val] = 1

        # Add current node to the ancestors list and recurse
        ancestors.append(node.val)
        self.dfs(node.left, matrix, ancestors)
        self.dfs(node.right, matrix, ancestors)

        # Remove current node from ancestors list (backtracking)
        ancestors.pop()

# Example usage
if __name__ == "__main__":
    root = TreeNode(0)
    root.left = TreeNode(1)
    root.right = TreeNode(2)
    root.right.left = TreeNode(3)
    root.right.right = TreeNode(4)

    solution = Solution()
    n = 5
    ancestor_matrix = solution.constructAncestorMatrix(root, n)

    for row in ancestor_matrix:
        print(row)

Complexity

⏰ Time complexity: O(n^2). Each row requires traversal of the subtree, and marking ancestors for all nodes can approach a quadratic complexity depending on the shape of the binary tree.
🧺 Space complexity: O(n). For recursive calls (stack space) or tracking ancestors at any point.

Problem#

Examples#

Example 1#

Example 2#

Similar problem#

Solution#

Method 1 - Using DFS#

Approach#

Code#

Complexity#