时间戳

首页 > 解决方案 > 如何有效地将 X 模 Y 添加到 numpy 数组中的每个元素？

python - 如何有效地将 X 模 Y 添加到 numpy 数组中的每个元素？

问题描述

假设我有一个非常大的 1 到 180 之间的值数组，这些值位于 s 数组中uint8（可以高达 255）。我想为每个值添加 90（模 180）：

original = np.array([1, 2, 3, 170, 171, 172], dtype=np.uint8)
modified = (original + 90) % 180

不幸的是，这会产生不正确的结果，因为在添加 90 时，较大的数字会在第一步中溢出它们的 uint8 整数：170 + 90 = 260，它大于 255。

# (170 + 90) % 180 is 80, not 4 :(
array([91, 92, 93,  4,  5,  6], dtype=uint8)

我在对性能非常敏感的环境中运行，我的输入列表非常大。因此，我想避免将此数组转换为更大数据类型的惩罚，并且我想使用高效的操作（例如，避免循环遍历数组并单独处理每个值）。

我怎样才能做到这一点？

标签： pythonarraysnumpymodulo

解决方案

一个非常简单的选项，因为您正在处理uint8，所以只需提前计算数组中每个可能值的结果并使用它：

import numpy as np
original = np.array([1, 2, 3, 170, 171, 172], dtype=np.uint8)
value_map = ((np.arange(256) + 90) % 180).astype(np.uint8)
modified = value_map[original]
print(modified)
# [91 92 93 80 81 82]

这样做的好处是它不需要超过 256-element 的任何额外内存value_map，并且对于任何更大的数组，您也将节省大部分计算。

针对铸造运行时间基准：

import numpy as np

def add_val_mod_cast(a, val, mod):
    return ((a.astype(np.uint16) + val) % mod).astype(np.uint8)

def add_val_mod_map(a, val, mod):
    value_map = ((np.arange(256) + val) % mod).astype(np.uint8)
    return value_map[a]

np.random.seed(0)
a = np.random.randint(256, size=10_000_000).astype(np.uint8)
val = 90
mod = 180
print((add_val_mod_cast(a, val, mod) == add_val_mod_map(a, val, mod)).all())
# True
%timeit add_val_mod_cast(a, val, mod)
# 72.6 ms ± 2.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit add_val_mod_map(a, val, mod)
# 40.8 ms ± 606 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

推荐阅读