高效重排2D NumPy阵列

Question

假设我有一个2D NumPy数组：

x = np.random.rand(100, 100000)

并检索按列排序的索引(即，每一列独立于其他列进行排序，并返回索引)：

idx = np.argsort(x, axis=0) 

然后，对于每一列，我需要indices = 10，20，30，40，50的值首先是(该列的)前5行，然后是其他排序的值(而不是索引！)。

一种天真的方法可能是：

indices = np.array([10, 20, 30, 40, 50])
out = np.empty(x.shape, dtype=int64)

for col in range(x.shape[1]):
    # For each column, fill the first few rows with `indices`
    out[:indices.shape[0], col] = x[indices, col]  # Note that we want the values, not the indices

    # Then fill the rest of the rows in this column with the remaining sorted values excluding `indices`
    n = indices.shape[0]
    for row in range(indices.shape[0], x.shape[0]):
        if idx[row, col] not in indices:
            out[n, col] = x[row, col]  # Again, note that we want the value, not the index
            n += 1

Answer 1

方法#1

这里有一个基于previous post的，不需要idx -

xc = x.copy()
xc[indices] = (xc.min()-np.arange(len(indices),0,-1))[:,None]
out = np.take_along_axis(x,xc.argsort(0),axis=0)

方法#2

另一个使用idx的np.isin掩码-

mask = np.isin(idx, indices)
p2 = np.take_along_axis(x,idx.T[~mask.T].reshape(x.shape[1],-1).T,axis=0)
out = np.vstack((x[indices],p2))

方法#2- Alternative如果您不断地编辑out以更改除那些数组之外的所有内容，则indices赋值可能是适合您的方法：

n = len(indices)
out[:n] = x[indices]

mask = np.isin(idx, indices)
lower = np.take_along_axis(x,idx.T[~mask.T].reshape(x.shape[1],-1).T,axis=0)
out[n:] = lower

Answer 2

这将通过消除最内部的循环和if条件来帮助您开始。首先，您可以将x[:, col]作为输入参数x传入。

def custom_ordering(x, idx, indices):
    # First get only the desired indices at the top
    out = x[indices, :]

    # delete `indices` from `idx` so `idx` doesn't have the values in `indices`
    idx2 = np.delete(idx, indices)

    # select `idx2` rows and concatenate
    out = np.concatenate((out, x[idx2, :]), axis=0)

    return out

Answer 3

以下是我对这个问题的解决方案：

rem_indices = [_ for _ in range(x.shape[0]) if _ not in indices]    # get all remaining indices
xs = np.take_along_axis(x, idx, axis = 0)                                        # the sorted array
out = np.empty(x.shape)

out[:indices.size, :] = xs[indices, :]                                                  # insert specific values at the beginning
out[indices.size:, :] = xs[rem_indices, :]                                         # insert the remaining values after the previous

如果我没弄错你的问题，请告诉我。