首页 > 解决方案 > 将排序后的数组拆分为带有子列表的列表

问题描述

我有一个 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?这是最高效的方式吗?

非常感谢!

标签: pythonpython-3.xnumpy

解决方案


方法#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

另一种更高效的方法是获取拆分索引,然后对数组进行切片,然后zipping-

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

基准测试

1000000对具有10000独特元素的元素进行矢量化方法的时间-

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

推荐阅读