python - 根据字符串元素的字符切片 Numpy 数组

Question

我在基于 bool 掩码的 Numpy 数组切片方面遇到了一些问题。

我可以成功地进行以下屏蔽，我在其中选择小于 10 的整数。

L1 = [1, 2, 3, 10, 20, 4]
arr = np.array(L1)
mask = arr[:] < 10
print(mask) # [ True  True  True False False  True]
print(arr[mask]) # [1 2 3 4] <-- CORRECT

同样的策略也适用于切片字符串数组以匹配特定字符串:

L2 = ['abc', 'bac', 'foo', 'az', 'bar', 'ac']
arr = np.array(L2)
mask = arr[:] == 'foo'
print(mask) # [False False  True False False False]
print(arr[mask]) # ['foo'] <-- CORRECT

但是，当检查数组中每个字符串的字符时，切片策略不起作用。在这里，我想选择数组中以字符“a”开头的字符串。

L2 = ['abc', 'bac', 'foo', 'az', 'bar', 'ac']
arr = np.array(L2)
mask = arr[:][0] == 'a'
print(mask) # False
print(arr[mask]) # [] <-- WRONG

如何正确创建该 mask ？

Answer 1

In [192]: alist = ['abc', 'bac', 'foo', 'az', 'bar', 'ac']
In [193]: arr = np.array(alist)

简单的列表理解:

In [194]: [a[0]=='a' for a in alist]
Out[194]: [True, False, False, True, False, True]

它也适用于数组，但速度较慢(数组迭代比列表慢):

In [195]: timeit [a[0]=='a' for a in alist]
707 ns ± 5.61 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [196]: timeit [a[0]=='a' for a in arr]
4.88 µs ± 9.44 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

String startswith 也是一个选项:

In [197]: [a.startswith('a') for a in alist]
Out[197]: [True, False, False, True, False, True]
In [198]: timeit [a.startswith('a') for a in alist]
1.14 µs ± 11.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

numpy 没有自己的字符串处理工具。它有np.char函数，但它们只是应用python字符串方法，没有速度提升:

In [200]: np.char.startswith(arr, 'a')
Out[200]: array([ True, False, False,  True, False,  True])
In [201]: timeit np.char.startswith(arr, 'a')
12.5 µs ± 297 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

如果您以列表开始和结束，您将获得最快的速度。迭代数组，或将 bool 列表转换回数组需要时间。

重新阅读您的代码，我发现您想要选择项目，而不仅仅是创建 mask 。那么列表理解应该是:

In [215]: [a for a in alist if a[0]=='a']
Out[215]: ['abc', 'az', 'ac']
In [216]: timeit [a for a in alist if a[0]=='a']
645 ns ± 3.36 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

====

至于你失败的代码

[:] 在任何表达式中都不会为您做任何事情:

In [213]: arr[:]
Out[213]: array(['abc', 'bac', 'foo', 'az', 'bar', 'ac'], dtype='<U3')

因此，您只是检查 arr 的第一个元素是否为“a”。您没有测试每个元素的第一个字符串。

In [214]: arr[0]=='a'
Out[214]: False