如何基于另一个NumPy数组对字符串NumPy数组元素进行分组

我有两个NumPy数组：

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

需要基于数组b中的子串对数组a的元素进行分组，最终得到如下结果：

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

不想用常规for循环实现，这里提供两种更高效的方案：

方案1：NumPy向量化操作实现 #

利用NumPy的广播特性和字符处理函数，批量完成匹配判断，再提取分组：

# 生成布尔匹配矩阵：每行对应b的一个子串，每列对应a的一个元素，值为该元素是否包含对应子串
match_mask = np.char.find(a, b[:, np.newaxis]) != -1
# 按行提取匹配的元素，转换为列表形式
c = [a[mask_row].tolist() for mask_row in match_mask]

核心的匹配逻辑是NumPy的向量化操作，比逐个循环判断效率提升明显，尤其适合大数据量场景。

方案2：借助pandas简化分组逻辑 #

如果项目中已经使用pandas，这种方式逻辑更直观：

import pandas as pd

# 将a转为pandas Series
s = pd.Series(a)
# 为每个元素找到对应的b中子串（假设每个a元素仅匹配一个b中的子串）
group_keys = s.apply(lambda x: next(sub for sub in b if sub in x))
# 按分组键聚合，转换为目标格式
c = [group.tolist() for _, group in s.groupby(group_keys)]

这个方案代码可读性更强，后续如果需要扩展分组规则也更方便。

注意事项 #

如果a中存在不匹配b任何子串的元素，上述方案会生成空列表（方案1）或抛出异常（方案2），可根据实际需求添加默认值处理逻辑。

内容的提问来源于stack exchange，提问作者user16613865