leetcode-go/medium/341.flatten-nested-list-iterator.js

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * function NestedInteger() {
 *
 *     Return true if this NestedInteger holds a single integer, rather than a nested list.
 *     @return {boolean}
 *     this.isInteger = function() {
 *         ...
 *     };
 *
 *     Return the single integer that this NestedInteger holds, if it holds a single integer
 *     Return null if this NestedInteger holds a nested list
 *     @return {integer}
 *     this.getInteger = function() {
 *         ...
 *     };
 *
 *     Return the nested list that this NestedInteger holds, if it holds a nested list
 *     Return null if this NestedInteger holds a single integer
 *     @return {NestedInteger[]}
 *     this.getList = function() {
 *         ...
 *     };
 * };
 */

Array.prototype.peek = function() {
    return this[this.length - 1]
}

/**
 * @constructor
 * @param {NestedInteger[]} nestedList
 */
var NestedIterator = function(nestedList) {
	this.st = [nestedList]
	this.idx = [0]
	this.findNext()
}

/**
 * @this NestedIterator
 * @returns {boolean}
 */
NestedIterator.prototype.hasNext = function() {
	return this.idx.length !== 0
}

/**
 * @this NestedIterator
 * @returns {integer}
 */
NestedIterator.prototype.next = function() {
	let val = this.st.peek()[this.idx.peek()].getInteger()
	this.idx[this.idx.length-1]++
	this.findNext()
	return val
}

NestedIterator.prototype.findNext = function() {
	while (this.idx.length !== 0) {
		if (this.idx.peek() == this.st.peek().length) { // Out of bound
			this.idx.pop()
			this.st.pop()
			if (this.st.length !== 0) {
				this.idx[this.idx.length-1]++
			}
		} else if (!this.st.peek()[this.idx.peek()].isInteger()) { // Go into list
			this.st.push(this.st.peek()[this.idx.peek()].getList())
			this.idx.push(0)
		} else { // Found
			break
		}
	}
}

/**
 * Your NestedIterator will be called like this:
 * var i = new NestedIterator(nestedList), a = [];
 * while (i.hasNext()) a.push(i.next());
*/