1522. Diameter of N-Ary Tree

Using variable to maintain 2 top big variables

T: O(n)

S: O(n)

/*
// Definition for a Node.
class Node {
    public int val;
    public List<Node> children;

    
    public Node() {
        children = new ArrayList<Node>();
    }
    
    public Node(int _val) {
        val = _val;
        children = new ArrayList<Node>();
    }
    
    public Node(int _val,ArrayList<Node> _children) {
        val = _val;
        children = _children;
    }
};
*/

class Solution {
    public int diameter(Node root) {
        int[] result = new int[1];
        dfs(root, result);
        return result[0];
    }
    private int dfs(Node node, int[] result) {
        if (node == null) {
            return 0;
        }
        int firstBig = 0;
        int secondBig = 0;
        for (Node next : node.children) {
            int cur = dfs(next, result);
            if (cur > firstBig) {
                secondBig = firstBig;
                firstBig = cur;
            } else if (cur > secondBig) {
                secondBig = cur;
            }
        }
        result[0] = Math.max(result[0], firstBig + secondBig);
        return Math.max(firstBig, secondBig)+1;
    }
}

Using PQ, also O(n), because only top 2 (first big and second big) to maintain in pq.

/*
// Definition for a Node.
class Node {
    public int val;
    public List<Node> children;

    
    public Node() {
        children = new ArrayList<Node>();
    }
    
    public Node(int _val) {
        val = _val;
        children = new ArrayList<Node>();
    }
    
    public Node(int _val,ArrayList<Node> _children) {
        val = _val;
        children = _children;
    }
};
*/

class Solution {
    public int diameter(Node root) {
        int[] result = new int[1];
        dfs(root, result);
        return result[0];
    }
    private int dfs(Node node, int[] result) {
        if (node == null) {
            return 0;
        }
        Queue<Integer> pq = new PriorityQueue<>();
        for (Node next : node.children) {
            pq.offer(dfs(next, result));
            while (pq.size() > 2) { // always maintain top 2 pq
                pq.poll();
            }
        }
        int secondBig = !pq.isEmpty() ? pq.poll() : 0;
        int firstBig = !pq.isEmpty() ? pq.poll() : 0;
        result[0] = Math.max(result[0], firstBig + secondBig);
        return Math.max(firstBig, secondBig)+1;
    }
}

Last updated