Определить блоки в отсортированном массиве целых чисел

Question

У меня есть отсортированный целочисленный массив, например, [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, конечно, медленное.

Есть подсказки?

Answer 1

Вот один с маскировкой и нарезкой -

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

Answer 2

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)