针对lexsort的2D数组广播1D数组:当考虑另一个向量时,用于独立地对每列进行排序的排列

Question

考虑数组 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与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')

和

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

---

题

有没有办法"广播"使用分组数组g在每一列中使用lexsort？有什么更有效的方法呢？

Answer 1

这是一种方法 -

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