Selecting elements in numpy array using regular expressions

One may select elements in numpy arrays as follows

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

This property is used by the np.where function to retrive indices:

indices = np.where(a>0.9)

What I would like to do is to be able to use regular expressions in such element selection. For example, if I want to select elements from b above that match the [Aab] regexp, I ne开发者_如何学Goed to write the following code:

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

This looks too verbouse for me. Is there any shorter and more elegant way to do this?

There's some setup involved here, but unless numpy has some kind of direct support for regular expressions that I don't know about, then this is the most "numpytonic" solution. It tries to make iteration over the array more efficient than standard python iteration.

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)