使用正则表达式选择numpy数组中的元素

Question

我们可以通过以下方式选择numpy数组中的元素：

a = np.random.rand(100)
sel = a > 0.5 #select elements that are greater than 0.5
a[sel] = 0 #do something with the selection

b = np.array(list('abc abc abc'))
b[b==a] = 'A' #convert all the a's to A's

这个特性被np.where函数用来获取索引：

indices = np.where(a>0.9)

我想做的是能够在选择元素时使用正则表达式。例如，如果我想从上面的b中选择符合[Aab]这个正则表达式的元素，我需要写以下代码：

regexp = '[Ab]'
selection = np.array([bool(re.search(regexp, element)) for element in b])

这对我来说看起来太繁琐了。有没有更简洁、更优雅的方法来做到这一点呢？

Answer 1

这里需要一些设置，但除非numpy有我不知道的直接支持正则表达式的功能，否则这是最“numpy风格”的解决方案。它试图让对数组的遍历比标准的Python遍历更高效。

import numpy as np
import re

r = re.compile('[Ab]')
vmatch = np.vectorize(lambda x:bool(r.match(x)))

A = np.array(list('abc abc abc'))
sel = vmatch(A)