基于NumPy实现多操作向量化的高效方案问询

架构老林

2026-8-20

基于NumPy高效实现条件筛选与随机选择的方案需求

问题背景

希望基于NumPy高效实现一系列操作，替代现有循环较多的方案。

输入参数

X：形状为(n,2)的矩阵，第二列为分类变量（取值为0,1,...,M，M≤10）
Q：随机数矩阵，行数与X中的分类数相同，列数≤10
R：长度为p的字符数组
S：字符串数组，长度等于X中的分类数
j：从零开始的索引

其中维度n和p可能很大，满足1≤n,p≤10^7。

示例输入

import numpy as np

X = np.array([[2, 0], [4, 0], [5, 1], [5, 0], [3, 1]])  # 两个分类：0,1
Q = np.array([[3,7,13,11], [19,23,31,29]] )
R = np.array(['a', 'c', 'b', 'b', 'b', 'f'])
S = np.array(['aacbbc', 'cabbca'])  # 字符串数量与分类数一致
j = 3

需求目标

需要高效实现以下逻辑的结果：

result = [[X[k,0],np.random.choice(Q[X[k,1]])] for k in range(len(X)) if R[X[k,0]] == S[X[k,1]][j]]

示例结果为result = [[2, 7], [4, 3], [3, 29]]（随机选择部分结果可能不同）。

现有尝试

import numpy as np

X = np.array([[2, 0], [4, 0], [5, 1], [5, 0], [3, 1]])
Q = np.array([[3,7,13,11], [19,23,31,29]])
R = np.array(['a', 'c', 'b', 'b', 'b', 'f'])
S = np.array(['aacbbc', 'cabbca'])
j = 3

int_result = []  # 理想状态下，该中间结果应为二维NumPy数组
for cat in range(len(S)):
    ind1 = np.where(X[:,1] == cat)[0]
    tmp1 = np.where(R[X[ind1][:,0]] == S[cat][3])[0]
    tmp2 = X[ind1[tmp1]]
    int_result.extend(tmp2)

result = []  # 理想状态下，该结果应为二维NumPy数组
for i in range(len(int_result)):
    tmp3 = [int_result[i][0], np.random.choice(Q[int_result[i][1]])]
    result.append(tmp3)