python - 使用索引数组跨多个轴重新排序 numpy 数组的最快方法

Question

假设我有一些数组 A，其中 A.shape = (x0,...,xn,r)。

我想通过根据相应的索引数组 ind 在维度 {x0,...,xn} 中重新排序来“解读”A ，其中 ind.shape = (n,A.size) 。最后一个维度 r 的顺序未指定。

这是迄今为止我想出的最好方法，但我认为你可以做得更好!例如，我可以在不复制的情况下获得 A 的重新排序 View 吗？

import numpy as np
def make_fake(n=3,r=3):
    A = np.array([chr(ii+97) for ii in xrange(n**2)]
        ).repeat(r).reshape(n,n,r)
    ind = np.array([v.repeat(r).ravel() for v in np.mgrid[:n,:n]])
    return A,ind

def scramble(A,ind):
    order = np.random.permutation(A.size)
    ind_shuf = ind[:,order]
    A_shuf = A.flat[order].reshape(A.shape)
    return A_shuf,ind_shuf

def unscramble(A_shuf,ind_shuf):
    A = np.empty_like(A_shuf)
    for rr in xrange(A.shape[0]):
        for cc in xrange(A.shape[1]):
            A[rr,cc,:] = A_shuf.flat[
            (ind_shuf[0] == rr)*(ind_shuf[1] == cc)
            ]
    return A

示例:

 >>> AS,indS = scramble(*make_fake())
 >>> print AS,'\n'*2,indS
[[['e' 'a' 'i']
  ['a' 'c' 'f']
  ['i' 'f' 'i']]

 [['b' 'd' 'h']
  ['f' 'c' 'b']
  ['g' 'h' 'c']]

 [['g' 'd' 'b']
  ['e' 'h' 'd']
  ['a' 'g' 'e']]] 

[[1 0 2 0 0 1 2 1 2 0 1 2 1 0 0 2 2 0 2 1 0 1 2 1 0 2 1]
 [1 0 2 0 2 2 2 2 2 1 0 1 2 2 1 0 1 2 0 0 1 1 1 0 0 0 1]] 

 >>> AU = unscramble(AS,indS)
 >>> print AU

[[['a' 'a' 'a']
  ['b' 'b' 'b']
  ['c' 'c' 'c']]

 [['d' 'd' 'd']
  ['e' 'e' 'e']
  ['f' 'f' 'f']]

 [['g' 'g' 'g']
  ['h' 'h' 'h']
  ['i' 'i' 'i']]]

Answer 1

这是执行此操作的一种方法:

def unscramble(A_shuf,ind_shuf):
    order = np.lexsort(ind_shuf[::-1])
    return A_shuf.flat[order].reshape(A_shuf.shape)

本质上，您以 n 个索引的形式获得每个项目的排名，即 a = (0, 0)、b = (0, 1)、c = (0, 2)、d = (1, 0) ... 等等。如果您对排名进行排序，您将需要重新排序以将项目按升序排列。您可以使用 lexsort 或者使用 numpy.ravel_multi_index 来获取整数排名，并对整数排名应用 argsort。如果解释不清楚，请告诉我。