基于字符串元素的字符切片Numpy数组

Question

我在使用基于布尔蒙版的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

如何正确创建该掩码？

- stackoverflowuser2010

Answer 1

使用 numpy.char.startswith:

arr[np.char.startswith(arr, "a")]

输出：

array(['abc', 'az', 'ac'], dtype='<U3')

请注意，默认情况下，此操作使用第一个索引（即0）。您可以使用start参数作为索引：

arr[np.char.startswith(arr, "a", 1)]

输出：

array(['bac', 'bar'], dtype='<U3')

- Chris

Answer 2

你可以使用简单的List Comprehension：

In [3403]: L2 = ['abc', 'bac', 'foo', 'az', 'bar', 'ac']

In [3404]: arr = np.array(L2)  

In [3411]: res = np.array([i for i in arr if i.startswith('a')])

In [3412]: res
Out[3412]: array(['abc', 'az', 'ac'], dtype='<U3')

或者，如果你想使用mask：

使用np.char.startswith

In [3415]: mask = np.char.startswith(arr, 'a')
In [3417]: print(arr[mask])
['abc' 'az' 'ac']

- Mayank Porwal

---

谢谢。列表推导式是更直观的方法。 - stackoverflowuser2010

Answer 3

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)

字符串方法 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)

如果您从列表的开头和结尾开始，则可获得最佳速度。 对数组进行迭代或将布尔列表转换回数组需要时间。

重新阅读您的代码，我发现您想要选择项目，而不仅仅是创建掩码。那么列表推导式应该是：

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

- hpaulj

---

我认为使用列表推导式构建掩码是最通用的方法。使用像np.char.startswith()这样的专用函数会很难记住。 - stackoverflowuser2010

感谢您还解释了为什么在我的代码中 arr[:][0] == 'a' 不起作用。 - stackoverflowuser2010

Answer 4

从numpy 1.23.0开始，您将能够在数组中切片字符串。以下内容将正常工作，而无需进行任何副本：

L2 = np.array(['abc', 'bac', 'foo', 'az', 'bar', 'ac'])
mask = L2[:, None].view('U1')[:, 0] == 'a'
print(mask)      # True, False, False, True, False, True
print(arr[mask]) # 'abc', 'az', 'ac'

我正在编写一个更简单的方法，基于这个改变（这里），它将允许以下操作：

mask = np.char.slice_(L2, stop=1) == 'a'

再次强调，完全矢量化且无需复制。

- Mad Physicist