使用 numpy 的 argsort 时避免内存错误答案

【问题标题】：Avoiding memory error when using numpy's argsort使用 numpy 的 argsort 时避免内存错误
【发布时间】：2018-01-09 14:50:00
【问题描述】：

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

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,:])]

但我想知道为什么上面的解决方案不会导致相同的结果？

【问题讨论】：

你想做什么？
嗨，Ignacio，我正在尝试按照 x 的（第二个轴）的顺序对 y 进行排序。
你确定要y[c_sort] 吗？做一个小例子，看看这是否是你真正需要的。您可能希望进行高级索引，如此处所示 - stackoverflow.com/questions/46103044/…
亲爱的 Divakar，我完全确定我需要 y[c_sort]。我希望根据第二个轴对 whole ndarray 进行排序。
那么，哪个是正确的版本？我认为y[c_sort] 会产生不同的结果。如果循环的那个是正确的，那么之前链接的帖子应该可以解决你的问题，如果你想优化的话。

标签： python numpy

【解决方案1】：

在 cmets 中讨论后，loopy 版本似乎是正确的版本。所以，为了优化它，我们可以使用advanced-indexing。因此，给定 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]

advanced-indexing take_along_axis 的通用辅助函数可能很有用。

【讨论】：