[leetcode 341]练练手

题目概述

原题链接

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. Input: [[1,1],2,[1,1]] Output: [1,1,2,1,1]

解题思路

如果只是遍历一遍，用类似二叉树的中序遍历就好了。但是因为要生成迭代子，那么就得 考虑用栈来模拟递归了。

当然这个题目是比较简单的，压栈的是一个tuple－(list, index)。第一个是list，第二个指向 list的第几个item。

第一次通过的解法

比较慢。

class NestedIterator(object):
    def __init__(self, nestedList):
        """
        Initialize your data structure here.
        :type nestedList: List[NestedInteger]
        """
        self.q = [(nestedList, 0)]

    def next(self):
        """
        :rtype: int
        """

        val = None
        not_find = True
        while self.q and not_find:
            arr, num = self.q.pop()
            if num < len(arr):
                ele = arr[num]
                if isinstance(ele, int): # 如果提交需要改成　ele.isInteger()
                    val = ele            # ................  ele.getInteger()
                    not_find = False
                    self.q.append((arr, num + 1))
                else:
                    self.q.append((arr, num + 1))
                    self.q.append((ele, 0))  # 如果提交改成，self.q.append((ele.getList(),0))

        return None if not_find else val

    def hasNext(self):
        """
        :rtype: bool
        """
        has_next = False
        while not has_next and self.q:
            arr, num = self.q.pop()
            if num < len(arr):
                ele = arr[num]
                if isinstance(ele, int):
                    has_next = True
                    self.q.append((arr, num))
                else:
                    self.q.append((arr, num + 1))
                    self.q.append((ele, 0))

        return has_next

if __name__ == "__main__":
    nestedlist = [[1, 2], 3, [4, 5], 6]
    i, v = NestedIterator(nestedlist), []
    while i.hasNext(): v.append(i.next())
    print(v)  # [1, 2, 3, 4, 5, 6]

优化后的解法

之前的解法大概是140ms跑通了44个测试。而大多数的解法则在70ms左右。 怎么改进呢？ has_next和next方法本质上是一样的，可以只保留一个next就可以实现两个方法了。 那么我们可以提前跑next方法，缓存两个值。

以下是改进后的速度提升了一倍。

class NestedIterator2(object):
    def __init__(self, nestedList):
        """
        Initialize your data structure here.
        :type nestedList: List[NestedInteger]
        """
        self.q = [(nestedList, 0)]
        self.next_val = self.next0()
        if self.next_val is not None :
            self.has_next = True
        else:
            self.has_next = False

    def next(self):
        curr_val = self.next_val
        if self.has_next:
            self.next_val = self.next0()

        if self.next_val is not None:
            self.has_next = True
        else:
            self.has_next = False

        return curr_val

    def hasNext(self):
        """
        :rtype: bool
        """
        return self.has_next

    def next0(self):
        """
        :rtype: int
        """

        val = None
        not_find = True
        while self.q and not_find:
            arr, num = self.q.pop()
            if num < len(arr):
                ele = arr[num]
                if isinstance(ele, int):
                    val = ele
                    not_find = False
                    self.q.append((arr, num + 1))
                else:
                    self.q.append((arr, num + 1))
                    self.q.append((ele, 0))

        return None if not_find else val