Update flatten-nested-list-iterator.cpp

C++/flatten-nested-list-iterator.cpp

@@ -1,5 +1,5 @@
 // Time:  O(n)
-// Space: O(n)
+// Space: O(1)
 
 /**
  * // This is the interface that allows for creating nested lists.
@@ -19,6 +19,40 @@
  * };
  */
 class NestedIterator {
+public:
+    using IT = vector<NestedInteger>::iterator;
+    NestedIterator(vector<NestedInteger> &nestedList) {
+        nodes_.emplace(nestedList.begin(), nestedList.end());
+    }
+
+    int next() {
+        const auto result = nodes_.top().first->getInteger();
+        ++nodes_.top().first;
+        return result;
+    }
+
+    bool hasNext() {
+        while (!nodes_.empty()) {
+            auto& cur = nodes_.top();
+            if (cur.first == cur.second) {
+                nodes_.pop();
+            } else if (cur.first->isInteger()) {
+                return true;
+            } else {
+                auto& nestedList = cur.first->getList();
+                ++cur.first;
+                nodes_.emplace(nestedList.begin(), nestedList.end());
+            }
+        }
+        return false;
+    }
+private:
+    stack<pair<IT, IT>> nodes_;
+};
+
+// Time:  O(n)
+// Space: O(n)
+class NestedIterator2 {
 public:
     NestedIterator(vector<NestedInteger> &nestedList) {
         for (int i = static_cast<int>(nestedList.size()) - 1; i >= 0; --i) {