Python实现m×n数组0/1掩码：含极小值判定为0的规则处理

规则说明 #

• 原数组中值为0的位置保留0，其余非0值默认置为1
• 小数点后连续出现4位及以上0的极小值（等价于绝对值小于1e-4，即0.0001），统一判定为0，不转换为1

注：如果需要匹配问题给出的示例输出（科学计数法e-15量级的数输出为1），可将极小值阈值调整为更小的1e-16即可。

实现思路 #

优先推荐用numpy实现，处理大规模数组效率更高；如果不想引入第三方依赖，也可以用纯Python遍历实现。核心判断逻辑如下：

1. 元素等于0 → 输出0
2. 元素绝对值小于阈值 → 输出0
3. 其余情况 → 输出1

代码实现 #

1. Numpy版本（推荐） #

import numpy as np

def arr_to_mask(arr, threshold: float = 1e-4):
    """
    将m×n数组转换为0/1掩码
    :param arr: 输入m×n数组
    :param threshold: 极小值判定阈值，默认1e-4（对应小数点后4个连续0）
    :return: 0/1掩码数组
    """
    np_arr = np.asarray(arr, dtype=np.float64)
    # 布尔矩阵转整数：满足非0且大于等于阈值的位置为1，其余为0
    mask = ((np_arr != 0) & (np.abs(np_arr) >= threshold)).astype(int)
    return mask.tolist()

2. 纯Python版本（无第三方依赖） #

def arr_to_mask_pure(arr, threshold: float = 1e-4):
    mask = []
    for row in arr:
        row_mask = []
        for num in row:
            if num == 0 or abs(num) < threshold:
                row_mask.append(0)
            else:
                row_mask.append(1)
        mask.append(row_mask)
    return mask

示例测试 #

# 问题给出的示例输入
input_arr = [
    [-2.21520694000000e-15, -1.18292704000000e-15, 5.42940708000000e-15],
    [-2.40108895000000e-15, 3.09784301000000e-15, -1.18292704000000e-14],
    [0, 0, 0],
    [1.5, 2.1, 1.4]
]

# 调整阈值为1e-16匹配示例输出
print(arr_to_mask(input_arr, threshold=1e-16))

输出结果和问题给出的示例完全一致。

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