Numpy 字符串数组：如何根据另一个数组中的元素对一个数组中的元素进行分组

Question

我有两个 numpy 数组：

a = np.array(['123_H2O_56', '345_CO2', 'H6C6_test', 'check_H2O', '67_H2O'])
b = np.array(['H2O', 'CO2', 'H6C6'])

我需要根据“b”numpy 数组中的值对“a”numpy 数组中的元素进行分组，并将结果保存到第三个 numpy 数组中。 我知道如何用 for 循环来做到这一点，但我想可能有更好的解决方案。 输出应如下所示：

c = [['123_H2O_56', 'check_H2O', '67_H2O'],
     ['345_CO2'],
     ['H6C6_test']]

Answer 1

使用 pandas 和正则表达式：

import re
import pandas as pd

a = np.array(['123_H2O_56', '345_CO2', 'H6C6_test', 'check_H2O', '67_H2O'])
b = np.array(['H2O', 'CO2', 'H6C6'])

pattern = '(%s)' % '|'.join(map(re.escape, b))
# '(H2O|CO2|H6C6)'

s = pd.Series(a)
out = s.groupby(s.str.extract(pattern, expand=False)).agg(list)

输出：

CO2                           [345_CO2]
H2O     [123_H2O_56, check_H2O, 67_H2O]
H6C6                        [H6C6_test]
dtype: object