Construct Tree from Ancestor Matrix, Java

Construct Tree from Ancestor Matrix

//'main' method must be in a class 'Rextester'.
//Compiler version 1.8.0_111

// https://www.geeksforgeeks.org/construct-tree-from-ancestor-matrix/

import java.util.*;
import java.lang.*;

class Rextester
{  
    public static void main(String[] args) {
        int[][] matrix = {
            { 0, 0, 0, 0, 0, 0 }, 
            { 1, 0, 0, 0, 1, 0 }, 
            { 0, 0, 0, 1, 0, 0 }, 
            { 0, 0, 0, 0, 0, 0 }, 
            { 0, 0, 0, 0, 0, 0 }, 
            { 1, 1, 1, 1, 1, 0 } 
        };
        Solution sol = new Solution();
        TreeNode root = sol.constructNaryTree(matrix);      
        System.out.println(root.toString());
    }
}

class Solution {
    public TreeNode constructNaryTree(int[][] matrix) {
        int n = matrix.length;
        if (n == 0) return null;
        int[] count = new int[n];
        for (int i = 0; i < n; i++)
            for (int j = 0; j < n; j++)
                if (matrix[i][j] == 1) count[i]++;
        
        Integer[] nodeId = new Integer[n];
        for (int i = 0; i < n; i++) nodeId[i] = i;
        Arrays.sort(nodeId, (a, b) -> count[a] - count[b]);

TreeNode[] nodes = new TreeNode[n];
        boolean[] parent = new boolean[n];
        for (int i: nodeId) {
            nodes[i] = new TreeNode(i);
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1 && !parent[j]) {
                    nodes[i].children.add(nodes[j]);
                    parent[j] = true;
                }
            }
        }
        return nodes[nodeId[n - 1]];
    }
}

class TreeNode {
    int val;
    List<TreeNode> children;
    
    TreeNode(int val) {
        this.val = val;
        children = new ArrayList<>();
    }

public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append('{').append(val).append(':');
        for (TreeNode node: children) sb.append(node.toString());
        sb.append('}');
        return sb.toString();
    }
}


run \| edit \| history \| help	0