向numpy数组添加元素的最佳方法

Question

我有一个numpy数组，我可以用append方法很简单地往里面添加一个元素，就像这样：

numpy.append(myarray, 1)

在这个例子中，我只是添加了一个整数 1。

但是，这样添加元素是最快的方法吗？我的数组非常长，有几万条数据。

或者，直接给数组指定位置赋值是不是更好呢？像这样：

myarray[123] = 1

Answer 1

如果你在程序运行结束时知道数组的大小，那么提前分配一个合适大小的数组会快很多，然后再设置里面的值。如果你需要动态添加元素，最好不要一次添加一个元素，而是尽量减少添加的次数，这样可以避免反复生成很多副本。你还可以对比一下使用 np.append、np.hstack 和 np.concatenate 等方法的时间差异，看看哪个更快。

Answer 2

往numpy数组里添加数据是非常低效的。这是因为每次添加数据时，程序都需要找到并分配整个数组的内存。根据不同的情况，有更好的方法可以使用。

如果你提前知道数组的长度，最好使用像 np.ones、np.zeros 或 np.empty 这样的函数来预先分配数组。

desired_length = 500
results = np.empty(desired_length)
for i in range(desired_length):
    results[i] = i**2

如果你不知道长度，可能更有效的做法是先把结果放在一个普通的列表里，等到最后再转换成数组。

results = []
while condition:
    a = do_stuff()
    results.append(a)
results = np.array(results)

---

这里是我电脑上的一些时间测试。

def pre_allocate():
    results = np.empty(5000)
    for i in range(5000):
        results[i] = i**2
    return results

def list_append():
    results = []
    for i in range(5000):
        results.append(i**2)
    return np.array(results)

def numpy_append():
    results = np.array([])
    for i in range(5000):
        np.append(results, i**2)
    return results

%timeit pre_allocate()
# 100 loops, best of 3: 2.42 ms per loop

%timeit list_append()
# 100 loops, best of 3: 2.5 ms per loop

%timeit numpy_append()
# 10 loops, best of 3: 48.4 ms per loop

所以你可以看到，预先分配和先用列表再转换的方式都要快得多。