python - 使用 numpy 的 argsort 时避免内存错误

Question

以下源代码在我的机器上产生内存错误:

import numpy as np
x = np.random.random([100,100,100])
y = np.random.random([100,100,100])
c_sort = np.argsort(x, axis = 2)
f = y[c_sort]

您有一个简单易用的方法来避免内存错误吗？

执行此操作的另一种方法是

x = np.random.random([100,100,100])
y = np.random.random([100,100,100])
f = np.zeros([100,100,100])
for i in range(100):
    for j in range(100):
        f[i,j,:] = y[i,j, np.argsort(x[i,j,:])]

但我想知道为什么上面的解决方案没有得到相同的结果？

Answer 1

经过评论中的讨论，看来循环版本是正确的版本。因此，为了优化它，我们可以使用高级索引。因此，给定 argsort 索引为 idx = np.argsort(x,axis=2)，我们可以得到 f 像这样 -

m,n = y.shape[:2]
f = y[np.arange(m)[:,None,None], np.arange(n)[:,None], idx]

高级索引的通用辅助函数 take_along_axis可能有用。