numpy - 对非常大的一维数组进行排序

Question

我即将第一次尝试 Pytables，我需要在每个时间步将数据写入 hdf 文件。我将有超过 100,000 个时间步长。完成后，我想按第 2 列对我的 100,000+ x 6 数组进行排序，即，我目前已按时间排序所有内容，但现在我需要按降雨率递减的顺序对数组进行排序（第 2 列）。我不确定如何从这里开始。我知道将整个数组放在内存中是不明智的。任何想法如何快速有效地做到这一点？

感谢任何建议。

Answer 1

我知道将整个数组放在内存中是不明智的。

你可能想多了。一个 100K x 6 的阵列float64只需要大约 5MB 的 RAM。在我的电脑上，对这样一个数组进行排序大约需要 27 毫秒：

In [37]: a = np.random.rand(100000, 6)

In [38]: %timeit a[a[:,1].argsort()]
10 loops, best of 3: 27.2 ms per loop

Answer 2

除非您有一台非常旧的计算机，否则您应该将整个数组放入内存中。假设它们是浮点数，它只需要100000*6*4./2**20 = 2.29Mb。双打的两倍。您可以使用 numpy 的 sort 或 argsort 进行排序。例如，您可以从第二列获取排序索引：

import numpy as np
a = np.random.normal(0, 1, size=(100000,6))
idx = a[:, 1].argsort()

然后使用这些索引您想要的列或整个数组：

b = a[idx]

您甚至可以使用不同类型的排序并检查它们的速度：

In [33]: %timeit idx = a[:, 1].argsort(kind='quicksort')
100 loops, best of 3: 12.6 ms per loop

In [34]: %timeit idx = a[:, 1].argsort(kind='mergesort')
100 loops, best of 3: 14.4 ms per loop

In [35]: %timeit idx = a[:, 1].argsort(kind='heapsort')
10 loops, best of 3: 21.4 ms per loop

因此，您会看到对于这种大小的数组来说，这并不重要。