如何使用Python获取指定长度的子列表？

Question

给定一个包含子列表的列表，我想提取出长度符合要求的子列表。如果子列表的长度小于指定的长度，那么就提取所有的内容。下面是具体的说明。

在下面的例子中，我要提取长度为2的子列表。如果子列表的长度大于2，我就只提取前两个元素，其他的就不管了。

输入：

A = [['A',[1,2,3]],['D',[3,4]],['E',[6,7]],['F',[1]],['G',[7,6,5,4]]]

输出：

B = [['A',[1,2]],['D',[3,4]],['E',[6,7]],['F',[1]],['G',[7,6]]]

我现在是这样做的，效果不错，但我在想有没有更简单的方法。

B=[]

for el in A:
  l = []

  if len(el[1]) > 2:
     l.append(el[0])
     l.append(el[1][0:2])
     B.append(l)

  else:
     l.append(el[0])
     l.append(el[1][0:2])
     B.append(l)

print B

Answer 1

如果你想用递归的方式来做

import collections

def doMagic(data):
    if isinstance(data, collections.Iterable):
        if all(type(elem) in [type(1), type('a')] for elem in data):
            return data[:2]
        else:
            return [doMagic(elem) for elem in data][:2]
    else:
        return data

if __name__ == '__main__':
    data = [['A',[1,2,3]],['D',[3,4]],['E',[6,7]],['F',[1]],['G',[7,6,5,4]]]
    print [doMagic(elem) for elem in data]

Answer 2

你可以使用 map 函数

>>> a=[[1,[1,2,3]],[2,[1]],[3,[1,2]]]
>>> map(lambda v: [v[0],v[1][:2]], a)
[[1, [1, 2]], [2, [1]], [3, [1, 2]]]

Answer 3

你可以使用列表推导式来实现这个：

In [71]: A = [['A',[1,2,3]],['D',[3,4]],['E',[6,7]],['F',[1]],['G',[7,6,5,4]]]

In [72]: [[x[0],x[1][:2]] for x in A]
Out[72]: [['A', [1, 2]], ['D', [3, 4]], ['E', [6, 7]], ['F', [1]], ['G', [7, 6]]]