Broadcast 1D array against 2D array for lexsort ：在考虑另一个向量时独立对每列进行排序的排列

考虑数组a

np.random.seed([3,1415])
a = np.random.randint(10, size=(5, 4))
a

array([[0, 2, 7, 3],
       [8, 7, 0, 6],
       [8, 6, 0, 2],
       [0, 4, 9, 7],
       [3, 2, 4, 3]])

我可以创造b其中包含对每列进行排序的排列。

b = a.argsort(0)
b

array([[0, 0, 1, 2],
       [3, 4, 2, 0],
       [4, 3, 4, 4],
       [1, 2, 0, 1],
       [2, 1, 3, 3]])

我可以排序a with b

a[b, np.arange(a.shape[1])[None, :]]

array([[0, 2, 0, 2],
       [0, 2, 0, 3],
       [3, 4, 4, 3],
       [8, 6, 7, 6],
       [8, 7, 9, 7]])

这是说明我正在寻找的输出的底漆。我想要一个数组b具有对相应列进行排序所需的排列a当还考虑lexsort与另一个数组。

np.random.seed([3,1415])
a = np.random.randint(10, size=(10, 4))
g = np.random.choice(list('abc'), 10)

a

array([[0, 2, 7, 3],
       [8, 7, 0, 6],
       [8, 6, 0, 2],
       [0, 4, 9, 7],
       [3, 2, 4, 3],
       [3, 6, 7, 7],
       [4, 5, 3, 7],
       [5, 9, 8, 7],
       [6, 4, 7, 6],
       [2, 6, 6, 5]])

g

array(['c', 'a', 'c', 'b', 'a', 'a', 'a', 'b', 'c', 'b'], 
      dtype='<U1')

我想生成一个数组b其中每一列都是必要的排列lexsort对应的列a。还有lexsort首先按定义的组对列进行排序g然后通过每列中的值a.

我可以通过以下方式生成结果：

r = np.column_stack([np.lexsort([a[:, i], g]) for i in range(a.shape[1])])
r

array([[4, 4, 1, 4],
       [5, 6, 6, 1],
       [6, 5, 4, 5],
       [1, 1, 5, 6],
       [3, 3, 9, 9],
       [9, 9, 7, 3],
       [7, 7, 3, 7],
       [0, 0, 2, 2],
       [8, 8, 0, 0],
       [2, 2, 8, 8]])

我们可以看到这是有效的

g[r]

array([['a', 'a', 'a', 'a'],
       ['a', 'a', 'a', 'a'],
       ['a', 'a', 'a', 'a'],
       ['a', 'a', 'a', 'a'],
       ['b', 'b', 'b', 'b'],
       ['b', 'b', 'b', 'b'],
       ['b', 'b', 'b', 'b'],
       ['c', 'c', 'c', 'c'],
       ['c', 'c', 'c', 'c'],
       ['c', 'c', 'c', 'c']], 
      dtype='<U1')

and

a[r, np.arange(a.shape[1])[None, :]]

array([[3, 2, 0, 3],
       [3, 5, 3, 6],
       [4, 6, 4, 7],
       [8, 7, 7, 7],
       [0, 4, 6, 5],
       [2, 6, 8, 7],
       [5, 9, 9, 7],
       [0, 2, 0, 2],
       [6, 4, 7, 3],
       [8, 6, 7, 6]])

Question

有没有办法“广播”分组数组的使用g用于每列lexsort？什么是更有效的方法来做到这一点？

这是一种方法 -

def app1(a, g):
    m,n = a.shape

    g_idx = np.unique(g, return_inverse=1)[1]
    N = g_idx.max()+1

    g_idx2D = g_idx[:,None] + N*np.arange(n)
    r_out = np.lexsort([a.ravel('F'), g_idx2D.ravel('F')]).reshape(-1,m).T
    r_out -= m*np.arange(n)
    return r_out

这个想法很简单，我们创建一个2D整数版本的网格g字符串数组，然后通过一个障碍来偏移每一列，该障碍会限制lexsort在每一列中搜索。

现在，在性能方面，似乎对于大型数据集，lexsort本身就会成为瓶颈。对于我们的问题，我们只处理两列。所以，我们可以创建自己的自定义lexsort根据偏移量缩放第二列，偏移量是第一列的最大数字限制。其实现看起来像这样 -

def lexsort_twocols(A, B):
    S = A.max() - A.min() + 1
    return (B*S + A).argsort()

因此，将其纳入我们提出的方法中并优化g_idx2D，我们会有一个像这样的正式函数 -

def proposed_app(a, g):
    m,n = a.shape

    g_idx = np.unique(g, return_inverse=1)[1]
    N = g_idx.max()+1

    g_idx2D = (g_idx + N*np.arange(n)[:,None]).ravel()
    r_out = lexsort_twocols(a.ravel('F'), g_idx2D).reshape(-1,m).T    
    r_out -= m*np.arange(n)
    return r_out

运行时测试

原来的做法：

def org_app(a, g):
    return np.column_stack([np.lexsort([a[:, i], g]) for i in range(a.shape[1])])

时间安排 -

In [763]: a = np.random.randint(10, size=(20, 10000))
     ...: g = np.random.choice(list('abcdefgh'), 20)
     ...: 

In [764]: %timeit org_app(a,g)
10 loops, best of 3: 27.7 ms per loop

In [765]: %timeit app1(a,g)
10 loops, best of 3: 25.4 ms per loop

In [766]: %timeit proposed_app(a,g)
100 loops, best of 3: 5.93 ms per loop