python - 确定排序的 numpy 整数数组中的 block

Question

我有一个排序整数数组，例如，[0, 0, 1, 1, 1, 2, 4, 4]，我想确定整数 block 从哪里开始以及 block 有多长。 block 大小很小，但数组本身可能非常大，因此效率很重要。 block 的总数也是已知的。

numpy.unique 可以解决这个问题:

import numpy


a = numpy.array([0, 0, 1, 1, 1, 2, 4, 4])
num_blocks = 4
print(a)

_, idx_start, count = numpy.unique(a, return_index=True, return_counts=True)

print(idx_start)
print(count)

[0 0 1 1 1 2 4 4]
[0 2 5 6]
[2 3 1 2]

但是很慢。我假设，鉴于输入数组的特定结构，存在更有效的解决方案。

例如，像这样简单的东西

import numpy

a = numpy.array([0, 0, 1, 1, 1, 2, 3, 3])
num_blocks = 4


k = 0
z = a[k]
block_idx = 0
counts = numpy.empty(num_blocks, dtype=int)
count = 0
while k < len(a):
    if z == a[k]:
        count += 1
    else:
        z = a[k]
        counts[block_idx] = count
        count = 1
        block_idx += 1
    k += 1
counts[block_idx] = count

print(counts)

给出 block 大小，一个简单的 numpy.cumsum 将给出 index_start。使用 Python 循环当然很慢。

有什么提示吗？

Answer 1

这是一个带有一些 mask 和切片的 -

def grp_start_len(a):
    m = np.r_[True,a[:-1] != a[1:],True] #np.concatenate for a bit more boost
    idx = np.flatnonzero(m)
    return idx[:-1], np.diff(idx)

sample 运行-

In [18]: a
Out[18]: array([0, 0, 1, 1, 1, 2, 4, 4])

In [19]: grp_start_len(a)
Out[19]: (array([0, 2, 5, 6]), array([2, 3, 1, 2]))

时间(来自@AGN Gazer 的解决方案的设置)-

In [24]: np.random.seed(0)

In [25]: a = np.sort(np.random.randint(1, 10000, 10000))

In [26]: %timeit _, idx_start, count = np.unique(a, return_index=True, return_counts=True)
1000 loops, best of 3: 411 µs per loop

# @AGN Gazer's solution
In [27]: %timeit st = np.where(np.ediff1d(a, a[-1] + 1, a[0] + 1))[0]; idx = st[:-1]; cnt = np.ediff1d(st)
10000 loops, best of 3: 81.2 µs per loop

In [28]: %timeit grp_start_len(a)
10000 loops, best of 3: 60.1 µs per loop

增加尺寸 10x -

In [40]: np.random.seed(0)

In [41]: a = np.sort(np.random.randint(1, 100000, 100000))

In [42]: %timeit _, idx_start, count = np.unique(a, return_index=True, return_counts=True)
    ...: %timeit st = np.where(np.ediff1d(a, a[-1] + 1, a[0] + 1))[0]; idx = st[:-1]; cnt = np.ediff1d(st)
    ...: %timeit grp_start_len(a)
100 loops, best of 3: 5.34 ms per loop
1000 loops, best of 3: 792 µs per loop
1000 loops, best of 3: 463 µs per loop