Python无循环实现3D矩阵越界索引并返回默认值

问题背景 #

现有3D NumPy矩阵，通过GUI选择的索引存在越界风险，需要在索引越界时返回默认值1000。原循环实现效率低下，已基于List实现1D数组的无循环越界索引方案，需改造为支持set_3d[X,Y,Z]调用的3D版本，且全程使用向量化操作避免循环。

解决方案 #

利用NumPy的向量化特性，通过掩码过滤合法索引，直接对批量索引进行合法性校验：合法索引取原矩阵值，越界位置填充默认值，全程无显式循环，保证大数据量下的效率。

代码实现 #

import numpy as np

class Safe3DArray:
    def __init__(self, arr, default_value=1000):
        assert arr.ndim == 3, "输入必须是3D NumPy数组"
        self.arr = arr
        self.default = default_value
        self.shape = arr.shape

    def __getitem__(self, indices):
        # 强制要求三维索引格式
        if not isinstance(indices, tuple) or len(indices) != 3:
            raise IndexError("必须使用三维索引格式：[X,Y,Z]")
        
        x_idx, y_idx, z_idx = indices
        # 统一转换为NumPy数组，兼容标量、列表、数组类型的索引
        x_idx = np.asarray(x_idx)
        y_idx = np.asarray(y_idx)
        z_idx = np.asarray(z_idx)
        
        # 生成合法索引的掩码：每个维度都在有效范围内
        valid_mask = (
            (x_idx >= 0) & (x_idx < self.shape[0]) &
            (y_idx >= 0) & (y_idx < self.shape[1]) &
            (z_idx >= 0) & (z_idx < self.shape[2])
        )
        
        # 初始化结果数组为默认值
        result = np.full_like(x_idx, self.default, dtype=self.arr.dtype)
        # 合法索引位置填充原矩阵的值
        result[valid_mask] = self.arr[
            x_idx[valid_mask],
            y_idx[valid_mask],
            z_idx[valid_mask]
        ]
        
        # 对标量索引返回单个值，批量索引返回数组
        return result.item() if result.ndim == 0 else result

使用示例 #

# 测试代码
if __name__ == "__main__":
    # 创建3x3x3的测试矩阵
    test_3d = np.arange(27).reshape(3,3,3)
    safe_3d = Safe3DArray(test_3d)
    
    # 单个合法索引
    print(safe_3d[1, 1, 1])  # 输出：13
    # 单个越界索引
    print(safe_3d[3, 1, 1])  # 输出：1000
    # 批量混合合法/越界索引
    print(safe_3d[[1, 3, -1], [2, 0, 1], [0, 5, 2]])  # 输出：[11 1000  20]

关键特性 #

• 无循环高效处理：全程基于NumPy向量化操作，避免Python循环，大数据量下性能远超循环实现
• 多索引类型支持：兼容标量、列表、NumPy数组类型的索引输入
• 自动处理正负索引：负索引按NumPy规则判断是否越界，越界则返回默认值
• 类型兼容：结果数组的 dtype 与原矩阵保持一致，避免类型转换问题

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