代码之家 › 专栏 › 技术社区 › Nico Schlömer David Maze

确定排序的numpy整数数组中的块

numpy arrays python

Nico Schlömer David Maze · 技术社区 · 8 年前

我有一个分类整数数组,例如。, [0, 0, 1, 1, 1, 2, 4, 4] ,我想确定整数块的起始位置以及块的长度。块大小很小,但阵列本身可能很大,因此效率很重要。块的总数也是已知的。

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)

给出块大小,以及 numpy.cumsum 会给 index_start .使用Python循环当然很慢。

有什么提示吗?

2 回复 | 直到 8 年前

Divakar 8 年前

这是一个有掩蔽和切片的-

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)

样本运行-

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

AGN Gazer 8 年前

np.where(np.ediff1d(a, None, a[0]))[0]

如果你想得到答案中的第一个“0”,请在后面加一个非零的数字 a[0] :

np.where(np.ediff1d(a, None, a[0] + 1))[0]

编辑(块长度):

啊,刚刚注意到你还想得到块的长度。然后,修改上述代码:

st = np.where(np.ediff1d(a, a[-1] + 1, a[0] + 1))[0]
idx = st[:-1]
cnt = np.ediff1d(st)

然后

>>> print(idx)
[0 2 5 6]
>>> print(cnt)
[2 3 1 2]

编辑2(定时测试)

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

In [70]: %timeit _, idx_start, count = np.unique(a, return_index=True, return_counts=True)
240 Âµs Â± 7.44 Âµs per loop (mean Â± std. dev. of 7 runs, 1000 loops each)

In [71]: %timeit st = np.where(np.ediff1d(a, a[-1] + 1, a[0] + 1))[0]; idx = st[:-1]; cnt = np.ediff1d(st)
74.3 Âµs Â± 816 ns per loop (mean Â± std. dev. of 7 runs, 10000 loops each)