Construct binary tree from an ancestor matrix

Problem

Given an ancestor matrix mat[n][n], where mat[i][j] = 1 represents that i is an ancestor of j in some binary tree and mat[i][j] = 0 otherwise, construct a binary tree such that all its values are from 0 to n-1. The input for the program is valid, ensuring that a binary tree can be constructed. There can be multiple binary trees for the same input; the program will construct any one of them.

Examples

Example 1

Input: 
mat = [[0, 1, 1], 
	[0, 0, 0], 
	[0, 0, 0]]

Output: [0, 1, 2]
   0
  / \
 1   2


Explanation: 
Root = 0
Left child = 1
Right child = 2
Since `0` is the ancestor of both `1` and `2`, it becomes the root, while `1` and `2` become its children.

Example 2

Input: 
mat = [[0, 1, 0, 0], 
	[0, 0, 1, 1], 
	[0, 0, 0, 0], 
	[0, 0, 0, 0]]

Output: 


Explanation: 

Root = 0
Left subtree = 1 → [2 and 3 as children]
`0` is the ancestor of `1`, `1` is the ancestor of `2` and `3`. So, `1` becomes a child of `0` and `2`, `3` become children of `1`.

graph TD;
    0["0"]
    1["1"]
    2["2"]
    3["3"]
    0 --> 1
    0 ~~~ N1:::hidden
    1 --> 2
    1 --> 3

classDef hidden display:none

Solution

Method 1 - BFS

Here is the approach:

Key Observation: The ancestor matrix provides information about the parent-child relationships in a binary tree.
- If mat[i][j] = 1, then i is a candidate to be the parent of j.
- If a node has no ancestors (sum of column values), it must be the root of the binary tree.
Tree Construction:
- Determine the root node (node with no ancestors).
- Process each node and establish parent-child relationships based on the matrix.
- Construct the binary tree in order to satisfy the parent-child hierarchy indicated in the matrix.

Approach

Step 1: Calculate Ancestor Counts: Use the matrix to compute the number of ancestors for every node (sum of each column).
Step 2: Identify Root Node: A node whose ancestor count is 0 is the root of the tree.
Step 3: Node Hierarchy Construction:
- Create a mapping to insert children for every node.
- Process nodes sequentially to construct the tree relationships (parent-child).
- Assign left and right children explicitly.
Step 4: Build the Binary Tree: For any valid ancestor matrix, use a recursive or iterative approach to attach nodes as left or right children.

Code

Java

class Solution {
    static class TreeNode {
        int val;
        TreeNode left, right;

        TreeNode(int val) {
            this.val = val;
            left = right = null;
        }
    }

    public static TreeNode constructBinaryTree(int[][] mat) {
        int n = mat.length;

        // Step 1: Calculate ancestor counts.
        int[] ancestorCount = new int[n];
        for (int col = 0; col < n; col++) {
            int count = 0;
            for (int row = 0; row < n; row++) {
                count += mat[row][col];
            }
            ancestorCount[col] = count;
        }

        // Step 2: Identify root node.
        Queue<Integer> queue = new LinkedList<>();
        TreeNode[] nodes = new TreeNode[n];
        for (int i = 0; i < n; i++) {
            if (ancestorCount[i] == 0) {
                queue.add(i);
            }
            nodes[i] = new TreeNode(i);
        }

        TreeNode root = null;

        // Step 3: Build tree.
        while (!queue.isEmpty()) {
            int parentIndex = queue.poll();
            
            if (root == null) root = nodes[parentIndex];

            // Attach children to the current parent.
            boolean createdLeft = false;
            for (int childIndex = 0; childIndex < n; childIndex++) {
                if (mat[parentIndex][childIndex] == 1 && --ancestorCount[childIndex] == 0) {
                    queue.add(childIndex);

                    if (!createdLeft) {
                        nodes[parentIndex].left = nodes[childIndex];
                        createdLeft = true;
                    } else {
                        nodes[parentIndex].right = nodes[childIndex];
                    }
                }
            }
        }
        return root;
    }

    public static void main(String[] args) {
        int[][] mat = {
            {0, 1, 1},
            {0, 0, 0},
            {0, 0, 0}
        };

        TreeNode root = constructBinaryTree(mat);
        System.out.println("Constructed binary tree root: " + root.val);
    }
}

Python

class Solution:
    def constructBinaryTree(self, mat):
        n = len(mat)

        # Step 1: Calculate ancestor counts.
        ancestor_count = [sum(row[col] for row in mat) for col in range(n)]

        # Step 2: Identify root node.
        queue = []
        nodes = [TreeNode(i) for i in range(n)]
        for i in range(n):
            if ancestor_count[i] == 0:
                queue.append(i)

        root = None

        # Step 3: Build the tree using BFS.
        while queue:
            parent_index = queue.pop(0)
            
            if root is None:
                root = nodes[parent_index]

            created_left = False  # Manage left-right child creation.
            for child_index in range(n):
                if mat[parent_index][child_index] == 1:
                    ancestor_count[child_index] -= 1
                    if ancestor_count[child_index] == 0:
                        queue.append(child_index)
                        if not created_left:
                            nodes[parent_index].left = nodes[child_index]
                            created_left = True
                        else:
                            nodes[parent_index].right = nodes[child_index]

        return root

# Example Usage:
mat = [
    [0, 1, 1],
    [0, 0, 0],
    [0, 0, 0]
]
solution = Solution()
root = solution.constructBinaryTree(mat)
print("Constructed binary tree root value:", root.val)

Complexity

⏰ Time complexity: O(n^2)
- Calculating ancestor counts: O(n^2) (Processing matrix rows and columns).
- Tree construction using BFS: O(n^2) (Traverse nodes and update ancestor counts).
- Overall: O(n^2).
🧺 Space complexity: O(n)
- TreeNode storage: O(n)
- Queue storage: O(n)
- Overall: O(n).