fd

Author: dingjikerbo

solution/src/main/java/com/inuker/solution/ZigzagIterator.java

@@ -15,7 +15,7 @@ public class ZigzagIterator {
 
     private List<Iterator<Integer>> mIterators = new ArrayList<>();
 
-    private int mCurIdx;
+    private int mIndex;
 
     public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
         init(v1, v2);
@@ -30,20 +30,18 @@ private void init(List<Integer>... lists) {
     }
 
     public int next() {
-        return mIterators.get(mCurIdx++).next();
+        mIndex %= mIterators.size();
+        Iterator<Integer> iterator = mIterators.get(mIndex);
+        int n = iterator.next();
+        if (!iterator.hasNext()) {
+            mIterators.remove(iterator);
+        } else {
+            mIndex++;
+        }
+        return n;
     }
 
     public boolean hasNext() {
-        while (!mIterators.isEmpty()) {
-            mCurIdx %= mIterators.size();
-
-            Iterator<Integer> cur = mIterators.get(mCurIdx);
-
-            if (cur.hasNext()) {
-                return true;
-            }
-            mIterators.remove(mCurIdx);
-        }
-        return false;
+        return !mIterators.isEmpty();
     }
 }

test/src/main/java/com/inuker/test/main.java

@@ -7,6 +7,7 @@
 import com.leetcode.library.TreeNode;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Iterator;
@@ -17,43 +18,43 @@
 public class main {
 
     public static void main(String[] args) {
+        ZigzagIterator iterator = new ZigzagIterator(
+                Arrays.asList(1,2,3),
+                Arrays.asList(4,5,6,7),
+                Arrays.asList(8,9));
+        while (iterator.hasNext()) {
+            System.out.print(iterator.next() + " ");
+        }
     }
 
-    class PeekingIterator implements Iterator<Integer> {
-
-        private Integer mNext;
+    public static class ZigzagIterator {
 
-        private Iterator<Integer> mIterator;
-
-        public PeekingIterator(Iterator<Integer> iterator) {
-            // initialize any member here.
-            mIterator = iterator;
-        }
+        private List<Iterator<Integer>> mIterators;
+        private int mIndex;
 
-        // Returns the next element in the iteration without advancing the iterator.
-        public Integer peek() {
-            if (mNext != null) {
-                return mNext;
+        public ZigzagIterator(List<Integer>... args) {
+            mIterators = new LinkedList<>();
+            for (List<Integer> list : args) {
+                if (!list.isEmpty()) {
+                    mIterators.add(list.iterator());
+                }
             }
-            mNext = mIterator.next();
-            return mNext;
         }
 
-        // hasNext() and next() should behave the same as in the Iterator interface.
-        // Override them if needed.
-        @Override
-        public Integer next() {
-            if (mNext != null) {
-                Integer ret = mNext;
-                mNext = null;
-                return ret;
+        public int next() {
+            mIndex %= mIterators.size();
+            Iterator<Integer> iterator = mIterators.get(mIndex);
+            int n = iterator.next();
+            if (!iterator.hasNext()) {
+                mIterators.remove(iterator);
+            } else {
+                mIndex++;
             }
-            return mIterator.next();
+            return n;
         }
 
-        @Override
         public boolean hasNext() {
-            return mNext != null || mIterator.hasNext();
+            return !mIterators.isEmpty();
         }
     }
 }