Эффективно заполните маску на основе стартовых показателей

Question

У меня есть 2D-массив (для этого примера, на самом деле это может быть ND), и я хотел бы создать для него маску, которая маскирует конец каждой строки. Например:

np.random.seed(0xBEEF)
a = np.random.randint(10, size=(5, 6))
mask_indices = np.argmax(a, axis=1)

Я хотел бы преобразовать mask_indices в логическую маску. В настоящее время я не могу придумать лучшего способа, чем

mask = np.zeros(a.shape, dtype=np.bool)
for r, m in enumerate(mask_indices):
    mask[r, m:] = True

Так что для

a = np.array([[6, 5, 0, 2, 1, 2],
              [8, 1, 3, 7, 1, 9],
              [8, 7, 6, 7, 3, 6],
              [2, 7, 0, 3, 1, 7],
              [5, 4, 0, 7, 6, 0]])

и

mask_indices = np.array([0, 5, 0, 1, 3])

я бы хотел увидеть

mask = np.array([[ True,  True,  True,  True,  True,  True],
                 [False, False, False, False, False,  True],
                 [ True,  True,  True,  True,  True,  True],
                 [False,  True,  True,  True,  True,  True],
                 [False, False, False,  True,  True,  True]])

Существует ли векторизованная форма этой операции?

В общем, я хотел бы иметь возможность делать это во всех измерениях, кроме того, которое определяет точки индекса.

Answer 1

I. Маскирование массива Ndim вдоль последней оси (строки)

Чтобы массив массива n-dim мог маскироваться вдоль строк, мы могли бы сделать -

def mask_from_start_indices(a, mask_indices):
    r = np.arange(a.shape[-1])
    return mask_indices[...,None]<=r

Пример выполнения -

In [177]: np.random.seed(0)
     ...: a = np.random.randint(10, size=(2, 2, 5))
     ...: mask_indices = np.argmax(a, axis=-1)

In [178]: a
Out[178]: 
array([[[5, 0, 3, 3, 7],
        [9, 3, 5, 2, 4]],

       [[7, 6, 8, 8, 1],
        [6, 7, 7, 8, 1]]])

In [179]: mask_indices
Out[179]: 
array([[4, 0],
       [2, 3]])

In [180]: mask_from_start_indices(a, mask_indices)
Out[180]: 
array([[[False, False, False, False,  True],
        [ True,  True,  True,  True,  True]],

       [[False, False,  True,  True,  True],
        [False, False, False,  True,  True]]])

II. Маскирование массива Ndim вдоль общей оси

Для маскирования массивов n-dim вдоль общей оси это будет -

def mask_from_start_indices_genericaxis(a, mask_indices, axis):
    r = np.arange(a.shape[axis]).reshape((-1,)+(1,)*(a.ndim-axis-1))
    mask_indices_nd = mask_indices.reshape(np.insert(mask_indices.shape,axis,1))
    return mask_indices_nd<=r

Примерные прогоны -

Настройка массива данных:

In [288]: np.random.seed(0)
     ...: a = np.random.randint(10, size=(2, 3, 5))

In [289]: a
Out[289]: 
array([[[5, 0, 3, 3, 7],
        [9, 3, 5, 2, 4],
        [7, 6, 8, 8, 1]],

       [[6, 7, 7, 8, 1],
        [5, 9, 8, 9, 4],
        [3, 0, 3, 5, 0]]])

Настройка и маскирование индексов по axis=1 -

In [290]: mask_indices = np.argmax(a, axis=1)

In [291]: mask_indices
Out[291]: 
array([[1, 2, 2, 2, 0],
       [0, 1, 1, 1, 1]])

In [292]: mask_from_start_indices_genericaxis(a, mask_indices, axis=1)
Out[292]: 
array([[[False, False, False, False,  True],
        [ True, False, False, False,  True],
        [ True,  True,  True,  True,  True]],

       [[ True, False, False, False, False],
        [ True,  True,  True,  True,  True],
        [ True,  True,  True,  True,  True]]])

Настройка и маскирование индексов по axis=2 -

In [293]: mask_indices = np.argmax(a, axis=2)

In [294]: mask_indices
Out[294]: 
array([[4, 0, 2],
       [3, 1, 3]])

In [295]: mask_from_start_indices_genericaxis(a, mask_indices, axis=2)
Out[295]: 
array([[[False, False, False, False,  True],
        [ True,  True,  True,  True,  True],
        [False, False,  True,  True,  True]],

       [[False, False, False,  True,  True],
        [False,  True,  True,  True,  True],
        [False, False, False,  True,  True]]])

---

Другие сценарии

A. Расширение до заданных конечных / стоп-индексов для маскировки

Чтобы расширить решения для случаев, когда нам даются конечные / стоп-индексы для маскирования, т.е. мы стремимся векторизовать mask[r, :m] = True, нам просто нужноотредактировать последний шаг сравнения в размещенных решениях следующим образом -

return mask_indices_nd>r

B. Вывод целочисленного массива

В некоторых случаях может потребоваться получить массив int. На тех, просто посмотрите на вывод как таковой. Следовательно, если out - это выход из опубликованных решений, то мы можем просто сделать out.view('i1') или out.view('u1') для int8 и uint8 dtype соответственно.

Для других типов данных мы бынеобходимо использовать .astype() для преобразования dtype.

C. Для маскировки с индексом для стоп-индексов

Для маскировки с индексом, то есть индекс должен быть включен для случая стоп-индексов, нам нужно просто включить равенство в сравнение. Следовательно, последний шаг будет -

return mask_indices_nd>=r

D. Для индексно-исключительной маскировки для начальных индексов

Это тот случай, когда заданы начальные индексы, и эти индексы не маскируются, а маскируются только со следующего элемента и далее до конца. Таким образом, аналогично рассуждениям, перечисленным в предыдущем разделе, для этого случая мы изменили бы последний шаг на -

return mask_indices_nd<r

Answer 2

>>> az = np.zeros(a.shape)
>>> az[np.arange(az.shape[0]), mask_indices] = 1
>>> az.cumsum(axis=1).astype(bool)  # use n-th dimension for nd case
array([[ True,  True,  True,  True,  True,  True],
       [False, False, False, False, False,  True],
       [ True,  True,  True,  True,  True,  True],
       [False,  True,  True,  True,  True,  True],
       [False, False, False,  True,  True,  True]])