Update Tree Diameter.java

varunu28 · web-flow · commit 1743d9da0de4 · 2024-02-25T16:55:42.000-08:00
diff --git a/Medium/Tree Diameter.java b/Medium/Tree Diameter.java
@@ -1,36 +1,39 @@
 class Solution {
-  public int treeDiameter(int[][] edges) {
-    Map<Integer, List<Integer>> graph = new HashMap<>();
-    for (int[] edge : edges) {
-      graph.computeIfAbsent(edge[0], k -> new ArrayList<>()).add(edge[1]);
-      graph.computeIfAbsent(edge[1], k -> new ArrayList<>()).add(edge[0]);
+    public int treeDiameter(int[][] edges) {
+        int n = edges.length;
+        List<Set<Integer>> graph = new ArrayList<>();
+        boolean[] visited = new boolean[n + 1];
+        for (int i = 0; i < n + 1; i++) {
+            graph.add(new HashSet<>());
+        }
+        for (int[] edge : edges) {
+            int nodeOne = edge[0];
+            int nodeTwo = edge[1];
+            graph.get(nodeOne).add(nodeTwo);
+            graph.get(nodeTwo).add(nodeOne);
+        }
+        int[] diameter = {0};
+        dfs(0, visited, graph, diameter);
+        return diameter[0];
     }
-    int farthestNode = bfsHelper(graph, 0)[0];
-    return bfsHelper(graph, farthestNode)[1];
-  }
 
-  private int[] bfsHelper(Map<Integer, List<Integer>> graph, int node) {
-    int farthestNode = 0;
-    int steps = 0;
-    Queue<Integer> queue = new LinkedList<>();
-    queue.add(node);
-    Set<Integer> visited = new HashSet<>();
-    while (!queue.isEmpty()) {
-      int size = queue.size();
-      while (size--  > 0) {
-        int removedNode = queue.remove();
-        visited.add(removedNode);
-        farthestNode = removedNode;
-        for (Integer peer : graph.getOrDefault(removedNode, new ArrayList<>())) {
-          if (!visited.contains(peer)) {
-            queue.add(peer);
-          }
+    private int dfs(int node, boolean[] visited, List<Set<Integer>> graph, int[] diameter) {
+        int highestDistance = 0;
+        int secondHighestDistance = 0;
+        visited[node] = true;
+        for (Integer conn : graph.get(node)) {
+            int distance = 0;
+            if (!visited[conn]) {
+                distance = 1 + dfs(conn, visited, graph, diameter);
+            }
+            if (distance > highestDistance) {
+                secondHighestDistance = highestDistance;
+                highestDistance = distance;
+            } else if (distance > secondHighestDistance) {
+                secondHighestDistance = distance;
+            }
         }
-      }
-      if (!queue.isEmpty()) {
-        steps++;
-      }
+        diameter[0] = Math.max(diameter[0], highestDistance + secondHighestDistance);
+        return highestDistance;
     }
-    return new int[]{farthestNode, steps};
-  }
 }
