Update Flatten Nested List Iterator.java

varunu28 · web-flow · commit b99d324a2f3e · 2023-10-20T07:50:48.000-07:00
diff --git a/Medium/Flatten Nested List Iterator.java b/Medium/Flatten Nested List Iterator.java
@@ -17,31 +17,34 @@
  */
 public class NestedIterator implements Iterator<Integer> {
 
-  private Stack<NestedInteger> stack;
-  
-  public NestedIterator(List<NestedInteger> nestedList) {
-    this.stack = new Stack<>();
-    for (int i = nestedList.size() - 1; i >= 0; i--) {
-      this.stack.push(nestedList.get(i));
+    private Stack<NestedInteger> stack;
+
+    public NestedIterator(List<NestedInteger> nestedList) {
+        this.stack = new Stack<>();
+        for (int i = nestedList.size() - 1; i >= 0; i--) {
+            this.stack.push(nestedList.get(i));
+        }
     }
-  }
 
-  @Override
-  public Integer next() {
-    int value = this.stack.pop().getInteger();
-    return value;
-  }
+    @Override
+    public Integer next() {
+        return stack.pop().getInteger();
+    }
 
-  @Override
-  public boolean hasNext() {
-    while (!this.stack.isEmpty() && !this.stack.peek().isInteger()) {
-      List<NestedInteger> ni = this.stack.pop().getList();
-      for (int i = ni.size() - 1; i >= 0; i--) {
-        this.stack.push(ni.get(i));
-      }
+    @Override
+    public boolean hasNext() {
+        while (!stack.isEmpty() && !stack.peek().isInteger()) {
+            NestedInteger popped = stack.pop();
+            if (popped == null) {
+                continue;
+            }
+            List<NestedInteger> list = popped.getList();
+            for (int i = list.size() - 1; i >= 0; i--) {
+                stack.push(list.get(i));
+            }
+        }
+        return !stack.isEmpty();
     }
-    return !this.stack.isEmpty();
-  }
 }
 
 /**
