умное назначение в двумерном массиве numpy на основе массива 1D

Question

У меня есть пустой массив 2D, и я хочу превратить его в значения -1 \ 1 на основе следующей логики:

а. найти argmax () каждой строки

б. на основе этого 1D массива (а) присвойте значения, которые он содержит значение 1

с. на основании отрицания этого 1D массива присваивают значение -1

Пример:

arr2D = np.random.randint(10,size=(3,3))
idx = np.argmax(arr2D, axis=1)

arr2D = [[5 4 1]
         [0 9 4]
         [4 2 6]]
idx = [0 1 2]

arr2D[idx] = 1
arr2D[~idx] = -1

что я получаю, это:

arr2D = [[-1 -1 -1]
         [-1 -1 -1]
         [-1 -1 -1]]

пока я хотел:

arr2D = [[1 -1 -1]
         [-1 1 -1]
         [-1 -1 1]]

ценю помощь, Спасибо

Answer 1

Подход № 1

Создайте маску с этими argmax -

mask = idx[:,None] == np.arange(arr2D.shape[1])

Затем используйте эти индексы, а затем используйте их для создания этих 1 и-1s массив -

out = 2*mask-1

В качестве альтернативы, мы могли бы использовать np.where -

out = np.where(mask,1,-1)

Подход № 2

Другой способ созданиямаска будет -

mask = np.zeros(arr2D.shape, dtype=bool)
mask[np.arange(len(idx)),idx] = 1

Затем получите out, используя один из методов, перечисленных в подходе № 1.

Подход № 3

Еще один способ был бы таким:

out = np.full(arr2D.shape, -1)
out[np.arange(len(idx)),idx] = 1

В качестве альтернативы, мы могли бы использовать np.put_along_axis для назначения -

np.put_along_axis(out,idx[:,None],1,axis=1)