基于Python/Numpy向量化（无迭代）实现指定规则的D列计算

需求说明 #

通过无迭代的向量化方式计算表格中的D列，规则对应示例如下：

• 当某行满足「A非零且B > C」时，向下查找第一个「C值 ≥ 该行B值」的行，将目标行的D列设为该行的A值
• 其余所有行的D列值为0

示例表格 #

向量化实现代码 #

import numpy as np

# 示例数据
A = np.array([1, 0, 0, 2, 0, 3, 0])
B = np.array([2, 0, 0, 4, 0, 5, 0])
C = np.array([0, 1, 2, 3, 4, 4, 6])

# 1. 筛选出需要匹配的源行索引（A非零且B > C）
source_mask = (A != 0) & (B > C)
source_indices = np.where(source_mask)[0]
source_B_values = B[source_indices]

# 2. 为每个源行找到对应的目标行索引（首个C ≥ 源行B且在源行之后的行）
# 创建广播式的匹配矩阵，标记所有C ≥ 源行B的位置
match_mask = C >= source_B_values[:, np.newaxis]
# 屏蔽源行及之前的位置，只保留源行之后的区域
match_mask[:, :source_indices[:, np.newaxis] + 1] = False

# 找到每个源行对应的首个匹配行索引
target_indices = np.argmax(match_mask, axis=1)
# 过滤掉没有找到匹配行的情况
valid_matches = np.any(match_mask, axis=1)
target_indices = target_indices[valid_matches]
source_indices = source_indices[valid_matches]

# 3. 生成D列
D = np.zeros_like(A)
D[target_indices] = A[source_indices]

print(D)  # 输出: [0 0 1 0 2 0 3]

代码说明 #

• 步骤1：通过布尔掩码筛选出需要处理的源行，即A非零且B大于C的行
• 步骤2：利用Numpy广播特性生成匹配矩阵，快速定位每个源行之后首个满足C≥源行B的行索引，全程无循环迭代
• 步骤3：初始化全0数组，将目标行的D值设为对应源行的A值

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