Перераспределение массива numpy с верхним пределом

Question

Сейчас я пытаюсь написать эффективную функцию, которая будет следовать простым правилам.

Предположим, матрица (она всегда сортировалась по убыванию).

| 7 4 3 2 |     # sums up to 16 
| 3 2 1 0 |     # sums up to 6

Я хочу обрезать значения в определенное число (пусть это будет 5), но «остаток» от отсечения должен распространяться на столбцы справа (эти столбцы также должны быть отсечены и т. д.). Таким образом, имея это в виду, матрица результатов должна выглядеть следующим образом:

| 5 5 4 2 |     # still sums up to 16, but values "moves" to the right
| 3 2 1 0 |     # no value extends 5 do nothing happens

Алгоритм не сложен в написании его одним l oop (сохранение отсеченных значений в буфер, распределение в следующем столбце и т. д.), что хорошо работает, но, как я уже сказал, было бы идеально сделать это с помощью векторизации (и отклонить с помощью любого l oop). Я попробовал некоторые решения cumsum + clip + diff, но на самом деле ничего не работает. Прямо сейчас я застрял.

Спасибо за любую помощь.

Answer 1

Если у вас много строк и не слишком много столбцов, вы можете подойти к нему следующим образом:

import numpy as np

def carry(sample,cap):
    result = sample.copy()
    for c in range(1,result.shape[1]):
        result[:,c] += np.maximum(result[:,c-1]-cap,0)
    result[:,:-1]  = np.minimum(result[:,:-1],cap)
    return result

вывод:

sample = np.array([[7, 4, 2,0], [3, 2, 1, 0]])
cap = 5

carry(sample,cap)

# [[5, 5, 3, 0],
   [3, 2, 1, 0]]

[РЕДАКТИРОВАТЬ] решение без l oop

Несмотря на то, что это не может использовать полную векторизацию (и работает медленнее), она делает трюк без каких-либо циклов:

def carry(sample,cap):
    fCarry = np.frompyfunc(lambda a,b:b+max(0,a-cap),2,1) 
    result = fCarry.accumulate(sample,dtype=np.object,axis=1)
    return np.minimum(cap,result.astype(sample.dtype))

Накопление с пользовательским ufun c переносит лишнюю сумму (через шапку) к следующему элементу. Затем все элементы переносятся в указанное ограничение, если они закончились (их перенос уже передан следующему соседу)