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

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

问题

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

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

Answer 2

根据 Divakar 的回答，我发布这个只是为了有一个展示我的衍生作品的好地方。他的 lexsort_twocols 函数可以完成我们需要的一切，并且可以很容易地应用于将一个维度广播到多个其他维度上。我们可以放弃 proposed_app 中的额外工作，因为我们可以在 lexsort_twocols 函数中的 argsort 中使用 axis=0。

def lexsort2(a, g):
    n, m = a.shape
    f = np.unique(g, return_inverse=1)[1] * (a.max() - a.min() + 1)
    return (f[:, None] + a).argsort(0)

lexsort2(a, g)

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

我也想到了这个...虽然没有那么好，因为我仍然依赖 np.lexsort，正如 Divakar 指出的那样，它可能很慢。

def lexsort3(a, g):
    n, m = a.shape
    a_ = a.ravel()
    g_ = np.repeat(g, m)
    c_ = np.tile(np.arange(m), n)
    return np.lexsort([c_, a_, g_]).reshape(n, m) // m

lexsort3(a, g)

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

假设我的第一个概念是 lexsort1

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

from timeit import timeit
import pandas as pd

results = pd.DataFrame(
    index=[100, 300, 1000, 3000, 10000, 30000, 100000, 300000, 1000000],
    columns=['lexsort1', 'lexsort2', 'lexsort3']
)

for i in results.index:
    a = np.random.randint(100, size=(i, 4))
    g = np.random.choice(list('abcdefghijklmn'), i)
    for f in results.columns:
        results.set_value(
            i, f,
            timeit('{}(a, g)'.format(f), 'from __main__ import a, g, {}'.format(f))
        )

results.plot()

再次感谢@Divakar。请为他的回答点赞！！！

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

Broadcast 1D array against 2D array for lexsort : Permutation for sorting each column independently when considering yet another vector

python

numpy

pandas

array-broadcasting