結構化排序性能 arrays (numpy)

Question

我有一個包含多個字段的數組，我想根據其中的 2 個字段對其進行排序。 這些字段之一是二進制的，例如：

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應該首先是 go。 這樣做（a） ：

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

與僅排序data['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% 的時間。 更改字段數據類型似乎也有一些影響 - 浮點數而不是整數似乎稍微快一些。

Answer 1

最簡單的方法是使用sort的order參數

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

但是，這在我的計算機上運行需要 121 毫秒，而data['class']和data['value']分別只需要 2.44 和 5.06 毫秒。 有趣的是， sort(data, order='class')再次花費了 135 毫秒，這表明問題在於對結構化數組進行排序。

因此，您使用argsort對每個字段進行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 毫秒。 不是很好，但如果性能至關重要，這是一種解決方法。

這似乎是一個已知問題： 對 numpy 結構化和記錄數組進行排序非常慢

編輯排序中使用的比較的源代碼可以在https://github.com/numpy/numpy/blob/dea85807c258ded3f75528cce2a444468de93bc1/numpy/core/src/multiarray/arraytypes.c.src 中看到。 數字類型要簡單得多。 盡管如此，性能的巨大差異還是令人驚訝。

Answer 2

除了@user2699 的良好（通用）答案之外，在您的特定情況下，您還可以作弊，因為結構化數組的兩個字段屬於相同的 integer 類型並且值相對較小（它們適合 32 位） . 作弊包括以下步驟：

• 使用arr - np.min(arr)
• 使用np.astype將每個字段轉換為np.uint64
• 使用以下命令將兩個字段打包成一個二進制數組： (class_arr << 32) | value_arr (class_arr << 32) | value_arr
• 使用np.sort對結果數組進行排序
• 使用以下命令解壓縮數組： class_arr = sorted_arr >> 32和value_arr = sorted_arr & ((1<<32)-1)

這種策略比使用兩個非常昂貴的np.argsort快得多。 對於更大的數組尤其如此，因為對大數組進行排序更加昂貴，並且np.sort比np.argsort更便宜。 更不用說間接索引在大數組上相對較慢，因為不可預測的偽隨機 memory 訪問模式和 RAM 的高延遲。 這種方法的缺點是實施起來有點棘手，而且它並不適用於所有情況。