向numpy数组批量添加元素的最优方案及NaN移除问题

问题场景 #

我需要在循环中向numpy数组添加约10万个元素，测试了两种实现方式：

import numpy as np
import time

# 方法1：循环调用np.append扩容
start = time.time()
b = np.array([0.0])
for i in range(1, 100000):
    b = np.append(b, np.array([i]))
end = time.time()
print(end-start)

# 方法2：预分配大数组后填充
start = time.time()
a = np.array([0])
A = np.empty(99999) * np.nan
a = np.concatenate((a, A), axis=0)
for i in range(1, 100000):
    a[i] = i
end = time.time()
print(end-start)

运行结果：

3.2555339336395264
0.018993854522705078

方法2的速度优势明显，但如果无法提前确定元素总数，预分配的数组会偏大，残留np.nan值，需要找到移除这些NaN的方法。

解决方法 #

1. 快速过滤NaN值 #

利用numpy的布尔索引结合np.isnan()，可以直接提取非NaN元素：

# 清理包含NaN的数组a
clean_array = a[~np.isnan(a)]

np.isnan(a)会生成一个布尔数组，标记每个位置是否为NaN；~符号对布尔值取反，得到所有有效元素的索引，通过切片即可得到清理后的数组。

2. 更高效的替代方案（避免NaN产生） #

如果预先无法确定元素数量，推荐先用Python列表动态添加元素，最后转换为numpy数组——这种方式效率和方法2接近，且完全不需要处理NaN：

start = time.time()
temp_list = [0.0]
for i in range(1, 100000):
    temp_list.append(i)
final_array = np.array(temp_list)
end = time.time()
print(end-start)

Python列表的append操作是均摊常数时间复杂度，配合最后一次转numpy数组，整体性能优异，代码也更简洁。

3. 动态扩容numpy数组（进阶） #

如果必须使用numpy数组进行动态添加，可以采用倍数扩容策略：当数组填满时，将容量翻倍，最后截断到实际元素数量，既避免np.append的频繁内存分配，也不会产生NaN：

start = time.time()
initial_capacity = 1000
arr = np.empty(initial_capacity, dtype=np.float64)
arr[0] = 0.0
current_size = 1

for i in range(1, 100000):
    if current_size == initial_capacity:
        # 扩容为当前容量的2倍
        initial_capacity *= 2
        arr = np.resize(arr, initial_capacity)
    arr[current_size] = i
    current_size += 1

# 截断到实际元素数量
final_arr = arr[:current_size]
end = time.time()
print(end-start)

这种方法在性能和灵活性之间取得了平衡，适合对numpy数组操作有强依赖的场景。

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