Python 基于Y==1条件填充二维numpy数组的实现问题

问题原因 #

你之前的写法是把Y中等于1的元素单独提取出来计算，丢失了原数组的位置信息，得到的结果形状和要求的(Y.shape[0], Y.shape[0])不匹配，所以不符合预期。

解决方法 #

我们可以先计算全量的原逻辑填充结果，再通过掩码过滤掉不需要修改的位置，保留初始值0即可，代码如下：

import numpy as np

# 你的输入数据
Y = np.array([[1], [2]])
# 初始化CN为全0
CN = np.zeros((Y.shape[0], Y.shape[0]), dtype=int)

# 第一步：计算原逻辑的全量填充结果
full_fill = np.where(Y.dot(Y.T) == np.square(Y), 1, 2)
# 第二步：构造掩码，满足以下任意条件的位置保留原填充结果，否则保持0
# 条件1：原逻辑填充值为2
# 条件2：原逻辑填充值为1，且该位置所在行对应的Y值为1
row_mask = (Y == 1).reshape(-1, 1) # 广播为(N,N)的行掩码
mask = (full_fill == 2) | ((full_fill == 1) & row_mask)

# 第三步：生成最终CN
CN = np.where(mask, full_fill, 0)

运行后得到的CN结果为：

[[1 2]
 [2 0]]

完全符合你的预期。

扩展说明 #

如果你的Y不是列向量，或者后续逻辑有调整，只需要修改row_mask的生成规则即可，不需要修改整体的计算逻辑。

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