Python:优雅高效的掩码列表方法

37

示例:

from __future__ import division
import numpy as np

n = 8
"""masking lists"""
lst = range(n)
print lst

# the mask (filter)
msk = [(el>3) and (el<=6) for el in lst]
print msk

# use of the mask
print [lst[i] for i in xrange(len(lst)) if msk[i]]

"""masking arrays"""
ary = np.arange(n)
print ary

# the mask (filter)
msk = (ary>3)&(ary<=6)
print msk

# use of the mask
print ary[msk]                          # very elegant  

结果如下:

>>> 
[0, 1, 2, 3, 4, 5, 6, 7]
[False, False, False, False, True, True, True, False]
[4, 5, 6]
[0 1 2 3 4 5 6 7]
[False False False False  True  True  True False]
[4 5 6]

正如您所看到的,对数组进行掩码操作比对列表更加优雅。如果您尝试在列表上使用数组掩码方案,您将会收到一个错误提示:

>>> lst[msk]
Traceback (most recent call last):
  File "<interactive input>", line 1, in <module>
TypeError: only integer arrays with one element can be converted to an index

这个问题是要找到一种优雅的掩码方式来处理列表。

更新:
jamylak 的答案因为介绍了compress而被接受,不过Joel Cornett提到的点让解决方案更加完整并符合我的兴趣。

>>> mlist = MaskableList
>>> mlist(lst)[msk]
>>> [4, 5, 6]
6个回答

62

如果你正在使用 numpy

>>> import numpy as np
>>> a = np.arange(8)
>>> mask = np.array([False, False, False, False, True, True, True, False], dtype=np.bool)
>>> a[mask]
array([4, 5, 6])

如果您没有使用numpy,那么您需要寻找itertools.compress

>>> from itertools import compress
>>> a = range(8)
>>> mask = [False, False, False, False, True, True, True, False]
>>> list(compress(a, mask))
[4, 5, 6]

2
迄今为止这里最好的解决方案 - Derek Eden
@jamylak:从Python 3+开始,似乎是使用zip()而不是izip() - Pierre
1
@Pierre 我已经更新了答案。我删除了那段代码片段,因为它可能会误导人,并且可以通过链接进行查看。另外,我认为由于原始问题使用了 numpy,所以突出显示 numpy 是很重要的。 - jamylak

16
如果您正在使用Numpy,您可以轻松地使用Numpy数组完成此操作,无需安装任何其他库:
>> a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>> msk = [ True, False, False,  True,  True,  True,  True, False, False, False]
>> a = np.array(a) # convert list to numpy array
>> result = a[msk] # mask a
>> result.tolist()
[0, 3, 4, 5, 6]

7

由于jamylak已经用实际的答案回答了这个问题,在这里我提供一个具有内置遮罩支持的列表示例(完全不必要):

from itertools import compress
class MaskableList(list):
    def __getitem__(self, index):
        try: return super(MaskableList, self).__getitem__(index)
        except TypeError: return MaskableList(compress(self, index))

使用方法:

>>> myList = MaskableList(range(10))
>>> myList
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
>>> mask = [0, 1, 1, 0]
>>> myList[mask]
[1, 2]

注意,compress函数在数据或掩码用尽时停止。如果您希望保留延伸到掩码长度以外的列表部分,可以尝试以下方法:
from itertools import izip_longest

[i[0] for i in izip_longest(myList, mask[:len(myList)], fillvalue=True) if i[1]]

+1 感谢您提出使用 MaskableList 来解决列表与数组上的掩码相同的问题。这看起来非常有趣,也符合我的需求。需要注意的是,与数组掩码相比,这些操作会稍微慢一些。我已将您的观点作为更新添加进去了。 - Developer
我尝试了你的“MaskableList”解决方案,但是在重新实例化时遇到了一些问题。在循环中的每个元素中,我想通过一个新列表来进行掩码:for i in arange(0,n): fts = MaskableList(F) sorter = argsort(A) result[i] = zip(fts[sorter],A[sorter])但是每次迭代,fts[sorter]包含相同的值,而sorter每次都不同。我通常将Python用作脚本语言,因此对对象不是很熟悉。 - Milla Well
@开发者:我没有特别测试过,但MaskableList可能会比较慢的一个原因是因为它正在进行稍微昂贵的异常处理。尝试切换try...except,使其默认尝试进行掩码处理。 - Joel Cornett
@MillaWell:我不熟悉argsort。另外,A是什么,F的内容是什么? - Joel Cornett
@MillaWell:啊,你的第一个问题是MaskableList并不像你想象的那样工作。它返回列表上二进制掩码(1、0或True/False)的结果。它不会根据索引列表重新排序元素。其次,zip(fts[sorter], A[sorter])将输出一个元组列表,但你有一个字典。 - Joel Cornett
显示剩余5条评论

4
我认为这不够优雅。虽然它很紧凑,但与大多数语言的结构非常不同,容易让人感到困惑。
正如Rossum所说,我们花费的时间更多地是在阅读代码而不是编写代码。如果一行代码的结构比较模糊,那么对于那些可能不熟悉Python但在其他语言方面具备完全能力的人来说,它就会变得更加令人困惑。
在服务代码的现实世界中,可读性比简短的符号表示更重要。就像修理汽车一样。有大量信息的大图可以使故障排除变得更加容易。
对我来说,我更愿意排除使用长形式的代码。
print [lst[i] for i in xrange(len(lst)) if msk[i]]

比起numpy的简短标记,它更容易理解。我不需要任何特定Python包的专业知识来解释它。


1

The following works perfectly well in Python 3:

np.array(lst)[msk]

如果您需要作为结果返回列表:
np.array(lst)[msk].tolist()

我来这里是为了解决这种方法在大型列表中效率低下的问题。哈哈 - Derek Eden

0

您也可以使用列表和zip函数

  1. 定义一个函数
def masklist(mylist,mymask):
    return [a for a,b in zip(mylist,mymask) if b]
  1. 使用它!
n = 8
lst = range(n)
msk = [(el>3) and (el<=6) for el in lst]
lst_msk = masklist(lst,msk)
print(lst_msk)

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接