Numpy.where on an array of strings using regex

user_11918648 0
Viewed 3319

I have an array like:

array = ['A0','A1','A2','A3','A4','B0','B1','C0']

and want to obtain an array which is true for values with an A followed by a number ranging from 0 to 2.

So far, this is the way I do it:

selection = np.where ((array == 'A0') | (array == 'A1') | (array == 'A2'), 1, 0)

But is there a more elegant way to do this by using e.g., a regular expresion like:

selection = np.where (array == 'A[0-1]', 1, 0)
numpy python
5 Answers

If using pandas is an option:

import numpy as np
import pandas as pd

a = np.array(['A0','A1','A2','A3','A4','B0','B1','C0'])
pd.Series(a).str.match(r'A[0-2]')
# 0     True
# 1     True
# 2     True
# 3    False
# 4    False
# 5    False
# 6    False
# 7    False
# dtype: bool
user_6366260

I don't think numpy if your best solution here. You can accomplish this using built-in python tools such as map.

import re

array = ['A0','A1','A2','A3','A4','B0','B1','C0']
p = r'A[0-2]'

list(map(lambda x: bool(re.match(p, x)), array))
# returns
[True, True, True, False, False, False, False, False]

# to get an array:
np.array(list(map(lambda x: bool(re.match(p, x)), array)))
# returns:
array([ True,  True,  True, False, False, False, False, False])
user_50037560

If it's not more complicated than A0, A1 and A2, you can use

a = np.array(['A0','A1','A2','A3','A4','B0','B1','C0'])
np.in1d(a, ['A0', 'A1', 'A2'])
# array([ True,  True,  True, False, False, False, False, False])
user_6366260

Try vectorization with blend of re :

import re
array = ['A0','A1','A2','A3','A4','B0','B1','C0']

y = np.vectorize(lambda y, x: bool(re.compile(x).match(y)))
selection = np.where(y(array, 'A[0-2]'), 1, 0)
print(selection)

#output:
[1 1 1 0 0 0 0 0]
user_91345280

You can also use list comprehension:

r = re.compile('A[0-2]')
selection = np.array([1 if re.match(r, i) else 0 for i in array])
user_51789050
Related