0341. Flatten Nested List Iterator (M)

Flatten Nested List Iterator (M)

题目

Given a nested list of integers, implement an iterator to flatten it.

Each element is either an integer, or a list -- whose elements may also be integers or other lists.

Example 1:

Input: [[1,1],2,[1,1]]
Output: [1,1,2,1,1]
Explanation: By calling next repeatedly until hasNext returns false, 
             the order of elements returned by next should be: [1,1,2,1,1].

Example 2:

Input: [1,[4,[6]]]
Output: [1,4,6]
Explanation: By calling next repeatedly until hasNext returns false, 
             the order of elements returned by next should be: [1,4,6].

---

题意

给定一个任意嵌套的整数列表，要求按顺序输出所有整数。

思路

迭代处理：用两个栈分别取存储暂时挂起的列表和对应的下标，每当遇到一个嵌套列表时，就把当前列表和下标压栈，进入嵌套列表继续处理；当当前列表已经遍历完时，则出栈之前的列表和下标继续处理。

也可直接用递归展开。

---

代码实现

Java

迭代处理

public class NestedIterator implements Iterator<Integer> {
    private Deque<List<NestedInteger>> lists = new ArrayDeque<>();
    private Deque<Integer> indices = new ArrayDeque<>();
    private List<NestedInteger> list;
    private int index;
    private Integer num;

    public NestedIterator(List<NestedInteger> nestedList) {
        list = nestedList;
        index = -1;
        generate();
    }

    @Override
    public Integer next() {
        Integer res = num;
        generate();
        return res;
    }

    @Override
    public boolean hasNext() {
        return num != null;
    }

    private void generate() {
        index++;

        while (index < list.size() || !lists.isEmpty()) {
            if (index == list.size()) {
                list = lists.pop();
                index = indices.pop() + 1;
            } else if (!list.get(index).isInteger()) {
                lists.push(list);
                indices.push(index);
                list = list.get(index).getList();
                index = 0;
            } else {
                break;
            }
        }

        num = index < list.size() ? list.get(index).getInteger() : null;
    }
}

递归展开

public class NestedIterator implements Iterator<Integer> {
    List<Integer> list = new ArrayList<>();
    int index = 0;

    public NestedIterator(List<NestedInteger> nestedList) {
        generate(nestedList);
    }

    @Override
    public Integer next() {
        return list.get(index++);
    }

    @Override
    public boolean hasNext() {
        return index != list.size();
    }

    private void generate(List<NestedInteger> nestedList) {
        for (NestedInteger item : nestedList) {
            if (item.isInteger()) {
                list.add(item.getInteger());
            } else {
                generate(item.getList());
            }
        }
    }
}

JavaScript

/**
 * @constructor
 * @param {NestedInteger[]} nestedList
 */
var NestedIterator = function (nestedList) {
  this.list = flatten(nestedList)
  this.index = 0
}

/**
 * @this NestedIterator
 * @returns {boolean}
 */
NestedIterator.prototype.hasNext = function () {
  return this.index < this.list.length
}

/**
 * @this NestedIterator
 * @returns {integer}
 */
NestedIterator.prototype.next = function () {
  return this.list[this.index++]
}

function flatten(nestedList) {
  let flattened = []
  for (const item of nestedList) {
    if (item.isInteger()) flattened.push(item.getInteger())
    else flattened = flattened.concat(flatten(item.getList()))
  }
  return flattened
}