Блог

Допустим, у нас есть два 3D-массива, A (x, y, z) и B (x, y, z), в которых x, y, z являются размерами. Я хочу определить все минимальные значения по оси z в массиве A, а затем на основе этих значений и их индексов выбрать соответствующие значения в B, сохранить их и заменить другие значения на ноль.

Вы можете думать об этом немного по-другому. Найти местоположения минимумов в A просто:

 ind = np.expand_dims(np.argmin(A, axis=2), axis=2)
  

Вы можете выполнить одно из следующих действий:

1. Самый простой: создайте замену B и заполните соответствующие элементы:

     C = np.zeros_like(B)
    np.put_along_axis(C, ind, np.take_along_axis(B, ind, 2), 2)
     

2. То же самое, но на месте:

     values = np.take_along_axis(B, ind, 2)
    B[:] = 0
    np.put_along_axis(B, ind, values, 2)
     

3. Преобразуйте индекс в маску:

     mask = np.ones(B.shape, dtype=bool)
    np.put_along_axis(mask, ind, False, 2)
    B[mask] = 0
     

Вы можете заменить вызовы take_along_axis и put_along_axis подходящими выражениями индексации. В частности:

 indx, indy = np.indices(A.shape[:-1])
indz = np.argmin(A, axis=-1)
  

Приведенные выше примеры затем преобразуются в

1. Новый массив:

     C = np.zeros_like(B)
    C[indx, indy, indz] = B[indx, indy, indz]
     

2. На месте:

     values = B[indx, indy, indz]
    B[:] = 0
    B[indx, indy, indz] = values
     

3. Маскированный:

     mask = np.ones(B.shape, dtype=bool)
    mask[indx, indy, indz] = False
    B[mask] = 0