6.1 元素级运算的内存账本

6.1 元素级运算的内存账本

本节摘要：元素级运算（加减乘除、比较、幂）对两个等形状数组执行时，账本上记三笔：读输入A、读输入B、分配并写输出C。每一次运算符调用都新分配输出，链式表达式会制造一串临时数组。算清这笔账是用好 out 参数（6.3 节）和估算程序内存峰值的基础。

三笔账模型

看最普通的加法：

import numpy as np

a = np.arange(1_000_000, dtype=np.float64)   # 约 8MB
b = np.arange(1_000_000, dtype=np.float64)   # 约 8MB

c = a + b
# 此刻内存里同时住着三块 8MB：a、b、c
# "c = a + b" 的真实动作：
# 1 分配一块新 8MB 给临时结果（三块并存）
# 2 逐元素计算写入临时块
# 3 把名字 c 绑到临时块
​

注意第三步：等号不复制数据，只是给结果起名字。所以 c 与那个临时块是同一块内存，没有第四次分配。账本定式：N 元运算 = N 次读 + 1 次分配写。

比较、幂、取模同理，都是三笔账：

import numpy as np

x = np.array([1.5, 2.5, 3.5])
y = np.array([0.5, 1.0, 1.5])

print(x + y)    # [2.  3.5 5. ]
print(x * y)    # [0.75 2.5  5.25]
print(x ** 2)   # [ 2.25  6.25 12.25]  一元运算：读一笔写一笔
print(x % 1)    # [0.5 0.5 0.5]
print(x > 2)    # [False  True  True]，输出 dtype 是 bool，每元素1字节
​

链式表达式：临时数组的连环账

表达式越长，临时数组越多。算一个常见的复合指标：

import numpy as np

a = np.random.rand(10_000_000)   # 80MB
b = np.random.rand(10_000_000)

z = np.sqrt(a ** 2 + b ** 2)
# 内部账单：
# t1 = a ** 2        分配 80MB
# t2 = b ** 2        分配 80MB
# t3 = t1 + t2       分配 80MB（此时 t1 t2 尚未释放）
# z  = sqrt(t3)      分配 80MB
# 峰值瞬间：a b t1 t2 t3 z 并存，约 480MB！
​

Python 会尽快释放不再引用的临时块，t1 在 t3 生成后即可回收，但峰值出现在 t1、t2、t3 同时在场的那一刻。数据规模翻十倍，这个表达式就可能把内存顶穿。

改进写法有两档。轻量档用 hypot 一类复合函数（一次调用零临时）：

z2 = np.hypot(a, b)              # 直接算平方和开根，不落中间数组
print(np.allclose(z, z2))        # True
​

重活档把链式拆开显式复用（配 6.3 节的 out 参数效果更佳）：

t = a * a                        # 就地复用思路见 6.3
t += b * b                       # 只有一个 80MB 临时
np.sqrt(t, out=t)                # 结果写回 t
print(np.allclose(z, t))         # True
​

运算的内存流向

结果 dtype：就高不就低

运算结果的类型由输入按提升规则决定，与精度直觉偶有冲突：

import numpy as np

i = np.array([1, 2, 3])            # int64
f = np.array([1.5, 2.5, 3.5])      # float64
print((i + f).dtype)               # float64，整型提升为浮点

small = np.array([100, 200], dtype=np.int16)
print((small + small).dtype)       # int16！int16 加 int16 不升 int32
# Windows 上 int32 加 int32 同理停在 int32，大数累加可能溢出

u = np.array([1, 2], dtype=np.uint8)
print(u - 2)   # [255 0]，无符号数下溢环绕
​

第 3 章的转型纪律在这里生效：对可能到边界的整数运算，提前 astype 到宽类型，一行代码换走一整类静默错误。

案例：温度序列的热指数

完整过程。背景：十万个 (温度, 湿度) 采样点，要计算简化热指数 HI = 0.8T + 0.02T×RH，评估两种写法的内存峰值。

操作：

import numpy as np

rng = np.random.default_rng(5)
T = rng.uniform(28, 40, 100_000)      # 摄氏度
RH = rng.uniform(30, 95, 100_000)     # 百分比

# 写法一：一步表达式，两个临时数组
hi1 = 0.8 * T + 0.02 * T * RH
# 账单：t1=0.8T、t2=0.02T、t3=t2*RH、t4=t1+t3，峰值约 4 份 0.8MB

# 写法二：复用中间量
hi2 = T * RH
hi2 *= 0.02
hi2 += 0.8 * T
print(np.allclose(hi1, hi2))          # True
# 账单：峰值约 2 份，砍掉一半

print("危险采样点占比:", (hi1 > 40).mean().round(4))   # 输出示例：0.0231
​

结果与解读：数值一致，写法二峰值内存减半；数据量大一个数量级时这就是"能不能跑起来"的差别。变式：把两列合成 (100000, 2) 的二维数组按列运算，能进一步统一管理多个派生指标。

💡 关键直觉：读代码时在脑子里记账——每个运算符一个临时块。账目清楚的人，能在写表达式那一刻就预估程序的内存峰值。

布尔运算的账本：输出缩水九成

有一个元素级运算的账本例外值得单独记账——比较运算的输出是 bool 数组，每元素只占一字节：

import numpy as np

a = np.random.rand(10_000_000)      # 80MB float64
flag = a > 0.5
print(flag.dtype, flag.nbytes / 1e6, "MB")   # bool 10.0 MB
# 输出只有输入的八分之一，比较运算的"写"比算术便宜得多

print(flag.sum())    # True 按 1 计数，统计不超过阈值外的元素个数
# 输出示例：5000723
​

这也是第 4 章掩码经济的底层原因：一次全量比较的产出只有八分之一大小，后续掩码取值才装得下。把"布尔很便宜"记进账本，写过滤管线时就不会误判成本。

本节要点回顾

• 三笔账定式：等形状二元运算 = 两笔读 + 一笔新分配写；等号只是命名，不复制
• 链式临时数组：表达式每个运算符制造一个临时块，峰值在它们同时在场时出现
• 复合函数优先：hypot 等内建复合函数一次完成，天然零中间量
• dtype 提升：整型同宽不升位，uint 下溢环绕；边界运算先转宽
• 估算习惯：内存峰值约等于 输入块 + 同时在场的临时块 + 输出块

下一节让形状不同的数组也能进账本：广播规则的四步对齐法。