python - 将排序数组拆分为带有子列表的列表

Question

我有一个经过排序的 float32 值数组，我想将此数组拆分为仅包含相同值的列表列表，如下所示:

>>> split_sorted(array) # [1., 1., 1., 2., 2., 3.]
>>> [[1., 1., 1.], [2., 2.], [3.]]

我目前的方法是这个函数

def split_sorted(array):
    split = [[array[0]]]

    s_index = 0
    a_index = 1
    while a_index < len(array):
        while a_index < len(array) and array[a_index] == split[s_index][0]:
            split[s_index].append(array[a_index])
            a_index += 1
        else:
            if a_index < len(array):
                s_index += 1
                a_index += 1
                split.append([array[a_index]])

我现在的问题是，有没有更 Pythonic 的方法来做到这一点？也许甚至可以用numpy？这是最高效的方式吗？

非常感谢!

Answer 1

方法#1

以a作为数组，我们可以使用np.split -

np.split(a,np.flatnonzero(a[:-1] != a[1:])+1)

示例运行 -

In [16]: a
Out[16]: array([1., 1., 1., 2., 2., 3.])

In [17]: np.split(a,np.flatnonzero(a[:-1] != a[1:])+1)
Out[17]: [array([1., 1., 1.]), array([2., 2.]), array([3.])]

方法#2

另一种更高效的方法是获取分割索引，然后对数组进行切片并压缩 -

idx = np.flatnonzero(np.r_[True, a[:-1] != a[1:], True])
out = [a[i:j] for i,j in zip(idx[:-1],idx[1:])]

方法#3

如果您必须获取子列表列表作为输出，我们可以通过列表复制来重新创建 -

mask = np.r_[True, a[:-1] != a[1:], True]
c = np.diff(np.flatnonzero(mask))
out = [[i]*j for i,j in zip(a[mask[:-1]],c)]

基准测试

对具有 10000 唯一元素的 1000000 元素进行矢量化方法的计时 -

In [145]: np.random.seed(0)
     ...: a = np.sort(np.random.randint(1,10000,(1000000)))

In [146]: x = a

# Approach #1 from this post
In [147]: %timeit np.split(a,np.flatnonzero(a[:-1] != a[1:])+1)
100 loops, best of 3: 10.5 ms per loop

# Approach #2 from this post
In [148]: %%timeit
     ...: idx = np.flatnonzero(np.r_[True, a[:-1] != a[1:], True])
     ...: out = [a[i:j] for i,j in zip(idx[:-1],idx[1:])]
100 loops, best of 3: 5.18 ms per loop

# Approach #3 from this post
In [197]: %%timeit
     ...: mask = np.r_[True, a[:-1] != a[1:], True]
     ...: c = np.diff(np.flatnonzero(mask))
     ...: out = [[i]*j for i,j in zip(a[mask[:-1]],c)]
100 loops, best of 3: 11.1 ms per loop

# @RafaelC's soln
In [149]: %%timeit
     ...: v,c = np.unique(x, return_counts=True)
     ...: out = [[a]*b for (a,b) in zip(v,c)]
10 loops, best of 3: 25.6 ms per loop