Numpy.where в массиве строк с использованием регулярных выражений

Question

У меня есть массив вроде:

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

и я хочу получить массив, который является истинным для значений с A, за которым следует число в диапазоне от 0 до 2.

Пока чтоВот как я это делаю:

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

Но есть ли более элегантный способ сделать это, используя, например, регулярное выражение типа:

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

Answer 1

При использовании панд это опция:

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

Answer 2

Если это не сложнее, чем A0, A1 и A2, вы можете использовать

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])

Answer 3

Я не думаю, что Numpy, если ваше лучшее решение здесь. Вы можете сделать это с помощью встроенных инструментов Python, таких как 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])

Answer 4

Вы также можете использовать list comprehension:

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

Answer 5

Попробуйте векторизацию со смесью 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]