как условно изменить значения массива

Question

Мне нужно добавить 2 новых столбца в существующий массив «условно» и хотел бы решить его без особых сложностей.

Предположим, что у меня есть следующий массив:

a = np.array([[1,  2], 
              [-1, 4],
              [1,  6],
              [-1, 8]])

К этому массиву мне нужно добавить 2 новых столбца, чтобы массив выглядел следующим образом:

a = np.array([[1,  2, 2, 0 ], 
              [-1, 4, 2, 4 ],
              [1,  6, 8, 4 ],
              [-1, 8, 8, 12]])

Вот логика:

Исходный массив состоит из 2 столбцов.И на выходе есть 4 столбца.(3. и 4. столбцы новые.)

Столбец 3. по сути является инкрементальной суммой второго столбца, но вы продолжаете добавлять значения только в столбец 2 (только из [:, 1])если соответствующее значение равно 1 в первом столбце массива (a [:, 0]).Например:

• a [0, 2] равно 2, потому что a [0, 1] равно 2, и мы принимаем его, потому что a [0,0] = 1
• a [1, 2] остается как "2", потому что a [1,0] = (- 1), так что мы пропускаем значение a [1, 1]
• a [2, 2] становится (2 +6 =) 8. Это сумма а [0, 2] и а [2, 2].Мы суммируем только значения из столбца 2. до тех пор, пока соответствующее значение строки НЕ (-1) в первом столбце.
• a [3, 2] остается равным 8, поскольку a [3, 0]= (-1), чтобы [3, 1] не добавлялось к сумме.

Создание столбца 4. такое же, но на этот раз вы просто добавляете значения из второго столбца.(a [:, 1]), пока значение строки из первого столбца равно -1.

Есть ли какая-нибудь библиотечная функция, поддерживающая такую ​​операцию хорошим способом?

Answer 1

Вот три строки:

a = np.array([[1,  2], 
              [-1, 4],
              [1,  6],
              [-1, 8]])

b = np.zeros(np.arange(2) + a.shape, a.dtype)
np.put_along_axis(b, *a[None].T, 1)
np.c_[a, b[:, 1:].cumsum(0)]
# array([[ 1,  2,  2,  0],
#        [-1,  4,  2,  4],
#        [ 1,  6,  8,  4],
#        [-1,  8,  8, 12]])

И вариант, который немного быстрее в основном потому, что он избавляется от удобного, но медленного np.c_

b = np.zeros_like(a)
np.put_along_axis(b, a[:,:1].clip(None, 0), a[:,1:], 1)
np.concatenate([a, b.cumsum(0)], 1)
# array([[ 1,  2,  2,  0],
#        [-1,  4,  2,  4],
#        [ 1,  6,  8,  4],
#        [-1,  8,  8, 12]])

Answer 2

Вы можете использовать np.cumsum для произведения маски в первом столбце во втором столбце, чтобы получить совокупную сумму, а затем numpy.stack , чтобы присоединиться к массивам:

import numpy as np

a = np.array([[1, 2],
              [-1, 4],
              [1, 6],
              [-1, 8]])

mask = a[:, 0] == 1
third = np.cumsum(a[:, 1] * mask).reshape(-1, 1)

mask = a[:, 0] == -1
fourth = np.cumsum(a[:, 1] * mask).reshape(-1, 1)

result = np.hstack((a, third, fourth))

print(result)

Вывод

[[ 1  2  2  0]
 [-1  4  2  4]
 [ 1  6  8  4]
 [-1  8  8 12]]

Обратите внимание, что необходимо изменить и third, и fourth, чтобы соответствовать размерам a.