Python NumPy：如何不转换更大dtype，限制数组加法至dtype最大值

问题背景 #

假设有两个uint8类型的NumPy数组：

arr1 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)
arr2 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)

我们需要实现饱和相加：当两数之和超过uint8的最大值255时，结果保留255，否则取实际和。预期结果为：

[100 200 255 255 255]

但直接执行数组加法会触发uint8的溢出特性（模256运算），得到不符合预期的结果：

import numpy as np

arr1 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)
arr2 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)

result = arr1 + arr2
print(result)

输出（后3个值因溢出错误）：

[100 200  44 144 244]

常规解决方案：转换为更大 dtype #

通过将数组转换为uint16（更大的无符号整数类型）计算，再截断回uint8，可得到正确结果：

import numpy as np

arr1 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)
arr2 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)

# 转换为uint16计算
sum_16 = arr1.astype(np.uint16) + arr2.astype(np.uint16)
# 截断超出255的值，转回uint8
result = np.where(sum_16 > 255, 255, sum_16).astype(np.uint8)
print(result)

输出符合预期：

[100 200 255 255 255]

更高效的无类型转换方法 #

如果不想转换 dtype，可以利用uint8的溢出特性直接判断，或者使用视图（view）避免数据复制，以下是两种高效方案：

方案1：利用溢出特性判断 #

uint8溢出时，两数之和会小于其中任意一个输入值（例如150+150=300，模256后为44，44<150）。基于这个特性，我们可以直接判断溢出并替换为255：

import numpy as np

arr1 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)
arr2 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)

sum_arr = arr1 + arr2
# 若sum_arr < arr1，说明发生溢出，设为255；否则保留原值
result = np.where(sum_arr < arr1, 255, sum_arr)
print(result)

输出：

[100 200 255 255 255]

这个方法完全不需要转换 dtype，仅通过一次加法和条件判断完成，性能最优。

方案2：使用视图（view）避免数据复制 #

view()方法不会复制数组数据，只是改变对内存的解释方式，比astype()更高效：

import numpy as np

arr1 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)
arr2 = np.array([50, 100, 150, 200, 250], dtype=np.uint8)

# 以int16视图读取uint8数组，相加后截断到0-255，再转回uint8
result = np.clip(arr1.view(np.int16) + arr2.view(np.int16), 0, 255).astype(np.uint8)
print(result)

输出同样符合预期，且因为view()不复制数据，比astype()的转换方案更快。

内容的提问来源于stack exchange，提问作者cdahms