python - 结构化数组排序的性能(numpy)

Question

我有一个包含多个字段的数组，我想根据其中两个字段进行排序。这些字段之一是二进制的，例如:

size = 100000
data = np.empty(
            shape=2 * size,
            dtype=[('class', int),
                   ('value', int),]
)

data['class'][:size] = 0
data['value'][:size] = (np.random.normal(size=size) * 10).astype(int)
data['class'][size:] = 1
data['value'][size:] = (np.random.normal(size=size, loc=0.5) * 10).astype(int)

np.random.shuffle(data)

我需要根据value对结果进行排序，对于相同的值，class=0应该先排序。这样做(a):

idx = np.argsort(data, order=['value', 'class'])
data_sorted = data[idx]

与仅对数据['value']进行排序相比，速度似乎慢了一个数量级。鉴于只有两个类，有没有办法提高速度？

通过随机实验，我注意到像这样的方法(b):

idx = np.argsort(data['value'])
data_sorted = data[idx]
idx = np.argsort(data_sorted, order=['value', 'class'], kind='mergesort')
data_sorted = data_sorted[idx]

比 (a) 花费的时间少约 20%。更改字段数据类型似乎也有一些效果 - float 而不是整数似乎稍微快一些。

Answer 1

最简单的方法是使用 sort 的 order 参数

sort(data, order=['value', 'class'])

但是，这在我的计算机上运行需要 121 毫秒，而 data['class'] 和 data['value'] 分别只需要 2.44 和 5.06 毫秒。有趣的是，sort(data, order='class') 再次花费了 135 毫秒，这表明问题出在结构化数组的排序上。

因此，您使用 argsort 对每个字段进行排序，然后对最终数组进行索引的方法似乎是正确的。但是，您需要单独对每个字段进行排序，

idx=argsort(data['class'])
data_sorted = data[idx][argsort(data['value'][idx], kind='stable')]

运行时间为 43.9 毫秒。 通过从索引中删除一个临时数组，您可以获得非常轻微的加速

idx = argsort(data['class'])
tmp = data[idx]
data_sorted = tmp[argsort(tmp['value'], kind='stable')]

运行时间为 40.8 毫秒。不太好，但如果性能至关重要的话，这是一个解决方法。

这似乎是一个已知问题: sorting numpy structured and record arrays is very slow

编辑 排序中使用的比较的源代码可以在 https://github.com/numpy/numpy/blob/dea85807c258ded3f75528cce2a444468de93bc1/numpy/core/src/multiarray/arraytypes.c.src 中看到。 。 数字类型要简单得多。尽管如此，性能上的巨大差异还是令人惊讶。