二维数组掩码处理及基于掩码索引的数组运算实现

嘿，我来帮你搞定这个把MATLAB逻辑转成Python的问题！先看看你原代码里的几个小问题：ma.where(x==0,x)不是正确创建掩码的方式，np.ones(x)应该改成np.ones_like(x)（因为x是数组，ones需要形状参数），还有y.not_masked这种写法是不对的。下面给你两种实用的解决方案，分别适配不同的后续需求：

方法一：用Numpy布尔索引（简单直接） #

这种方法最适合你当前的需求，逻辑清晰，代码简洁：

import numpy as np

def my_func(x):
    # 先把输入的列表转成Numpy数组，兼容列表/数组输入
    x_arr = np.asarray(x)
    # 创建和输入同形状的数组，全部初始化为π
    result = np.full_like(x_arr, np.pi, dtype=np.float64)
    # 找到所有非零元素的位置，把这些位置的元素平方后赋值回去
    non_zero_positions = x_arr != 0
    result[non_zero_positions] = x_arr[non_zero_positions] ** 2
    return result

# 测试一下
my_array = [[0,1,2],[1,0,2],[1,2,0]]
result_array = my_func(my_array)
print(result_array)

运行后你会得到期望的结果：

[[3.14159265 1.         4.        ]
 [1.         3.14159265 4.        ]
 [1.         4.         3.14159265]]

方法二：用掩码数组（适配后续除法等复杂运算） #

考虑到你提到后续函数会包含除法操作，且无法提前知晓索引位置，用Numpy的掩码数组（numpy.ma）会更稳妥——它会自动跟踪哪些元素是“被屏蔽”的，避免后续运算中出现除以零之类的错误：

import numpy as np
import numpy.ma as ma

def my_func_with_mask(x):
    x_arr = np.asarray(x)
    # 创建掩码数组：把所有值为0的元素标记为掩码（不参与后续运算）
    masked_x = ma.masked_equal(x_arr, 0)
    # 只对非掩码的元素执行平方运算
    squared_masked = masked_x ** 2
    # 把掩码的位置替换成π，得到最终数组
    result = squared_masked.filled(np.pi)
    return result

# 测试
my_array = [[0,1,2],[1,0,2],[1,2,0]]
result_array = my_func_with_mask(my_array)
print(result_array)

这个方法的优势在于，如果你后续要对结果数组做除法（比如1 / result_array），掩码数组会自动跳过原来的0位置（现在是π），不会让你意外触发除以零的错误，而且可以随时恢复掩码状态做其他统计运算。

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