基于三个2D NumPy数组的条件值替换问题求解

原始数组定义 #

import numpy as np

# 创建三个2x2数组
one = np.array([[1, -1], 
                [-1, 1]])

two = np.array([[-1, 0], 
                [-2, 2]])

three = np.array([[1, 1], 
                [-4, 3]])

需求说明 #

生成第四个2x2数组，规则如下：

• 若某个位置上三个数组的值都小于1，新数组对应位置设为-1
• 若某个位置上三个数组的值都大于1，新数组对应位置设为1
• 不满足上述两个条件的位置设为0

原代码问题分析 #

当前代码存在两个核心问题：

1. 初始数组错误：复制one数组作为final的初始值，导致未被条件覆盖的位置保留了one的原始值，而非需求的0
2. 条件判断错误：需求是判断值是否都大于1或都小于1，但代码中写的是>0和<0，不符合规则

正确实现代码 #

方法一：全0初始化后赋值 #

import numpy as np

one = np.array([[1, -1], 
                [-1, 1]])
two = np.array([[-1, 0], 
                [-2, 2]])
three = np.array([[1, 1], 
                [-4, 3]])

# 初始化与输入数组形状一致的全0数组
final = np.zeros_like(one)

# 三个值都大于1的位置设为1
final[(one > 1) & (two > 1) & (three > 1)] = 1

# 三个值都小于1的位置设为-1
final[(one < 1) & (two < 1) & (three < 1)] = -1

print(final)

方法二：嵌套np.where写法 #

import numpy as np

one = np.array([[1, -1], 
                [-1, 1]])
two = np.array([[-1, 0], 
                [-2, 2]])
three = np.array([[1, 1], 
                [-4, 3]])

# 嵌套判断：先判断全大于1，再判断全小于1，其余设0
final = np.where(
    (one > 1) & (two > 1) & (three > 1),
    1,
    np.where(
        (one < 1) & (two < 1) & (three < 1),
        -1,
        0
    )
)

print(final)

运行结果 #

两种方法都会输出期望的结果：

array([[ 0,  0],
       [-1,  1]])

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