python - 针对 lexsort : Permutation for sorting each column independently when considering yet another vector 的二维数组广播一维数组

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

我可以用 b 对 a 进行排序

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