时间戳

首页 > 解决方案 > 具有上限的numpy数组的重新分配

python - 具有上限的numpy数组的重新分配

问题描述

我目前正在努力编写遵循简单规则的高效函数。

让我们假设矩阵（它总是降序排序）。

| 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

算法在用一个循环编写它时并不难（将剪裁的值保存到缓冲区，将其分发到下一列，等等），效果很好，但正如我所说，使用矢量化来实现这一点是理想的（并使用任何循环关闭）。我尝试了一些 cumsum + clip + diff 解决方案，但没有任何效果。现在我被困住了。

感谢您的任何帮助。

标签： pythonnumpy

解决方案

如果你有很多行而不是太多列，你可以这样处理：

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]]

[编辑] 没有循环的解决方案

尽管这可能无法利用完全矢量化（并且运行速度较慢），但它可以在没有任何循环的情况下解决问题：

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))

使用自定义 ufunc 累积将额外的金额（超过上限）带到下一个元素。然后所有元素在结束时被降低到指定的上限（它们的进位已经转移到下一个邻居）

推荐阅读