[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。

第一次通过的解法 

比较慢。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
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方法,缓存两个值。

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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