python - 如何有条件地沿numpy ndarray的特定 Axis 在特定位置设置值

我正在寻找一种有效的Python方法，根据条件测试沿ndarray中的单个轴设置nth值。这是一个例子。

初始ndarray：

array([[[111, 112, 113, 114],
        [121, 122, 123, 124]],

       [[211, 212, 213, 214],
        [221, 222, 223, 224]],

       [[311, 312, 313, 314],
        [321, 322, 323, 324]]])

我想测试沿最内轴（轴2）在2位置的项目，换句话说，此示例中所有恰好以3结尾的数字。如果该数字低于截止值，我们将例如215，然后将该值设置为0。

结果如下所示：

array([[[111, 112, 0, 114],
        [121, 122, 0, 124]],

       [[211, 212, 0, 214],
        [221, 222, 223, 224]],

       [[311, 312, 313, 314],
        [321, 322, 323, 324]]])

我尝试了布尔索引，numpy.where等的组合，但是在使用花哨索引的同时沿特定轴隔离第n个项目似乎比我能找到的任何示例都要复杂。我希望这提出了一个有趣的问题，感谢您的帮助。

最佳答案

要访问沿轴0的所有第二列，请使用[:,:,2]（或a[...,2]）：

a[:,:,2]

array([[113, 123],
       [213, 223],
       [313, 323]])

要根据某些条件（例如

mask = a[:,:,2] < 215
a[:,:,2][mask] = 0

结果：

array([[[111, 112,   0, 114],
        [121, 122,   0, 124]],

       [[211, 212,   0, 214],
        [221, 222, 223, 224]],

       [[311, 312, 313, 314],
        [321, 322, 323, 324]]])