Быстрый способ извлечь подматрицу из матрицы numpy

Чтобы оптимизировать идею, нужно минимизировать вычисления, как только мы окажемся внутри цикла. Итак, имея это в виду, мы бы переставили строки массива, отсортированные по первому столбцу. Затем получите индексы, определяющие границы. Наконец, запустите наш цикл и просто нарежьте каждую группу, чтобы получить подматрицу на каждой итерации. Нарезка практически бесплатна при работе с массивами, так что это должно нам помочь.

Таким образом, одна реализация будет -

a0 = data_mat[:,0]
sidx = a0.argsort()
sd = data_mat[sidx] # sorted data_mat
idx = np.flatnonzero(np.concatenate(( [True], sd[1:,0] != sd[:-1,0], [True] )))
for i,j in zip(idx[:-1], idx[1:]):
    tmp_mat = sd[i:j]
    print tmp_mat

Если вы хотите сохранить каждую подматрицу в виде массива, чтобы иметь список массивов в качестве окончательного вывода, просто выполните -

[sd[i:j] for i,j in zip(idx[:-1], idx[1:])]

Для отсортированных data_mat

В случае, когда data_mat уже отсортировано, как показано в примере, мы могли бы избежать сортировки всего массива и напрямую использовать первый столбец, например:

a0 = data_mat[:,0]
idx = np.flatnonzero(np.concatenate(( [True], a0[1:] != a0[:-1], [True] )))

for i,j in zip(idx[:-1], idx[1:]):
    tmp_mat = data_mat[i:j]
    print(tmp_mat)

Опять же, чтобы получить все эти подматрицы в виде списка массивов, используйте -

[data_mat[i:j] for i,j in zip(idx[:-1], idx[1:])]

Обратите внимание, что подматрицы, которые мы получили бы с этим, будут в другом порядке, чем при сортировке, выполненной в предыдущем подходе.

Сравнительный анализ для отсортированных data_mat

Подходы -

# @Daniel F's soln-2
def split_app(data_mat):
    idx = np.flatnonzero(data_mat[1:, 0] != data_mat[:-1, 0]) + 1
    return np.split(data_mat, idx)

# Proposed in this post
def zip_app(data_mat):
    a0 = data_mat[:,0]
    idx = np.flatnonzero(np.concatenate(( [True], a0[1:] != a0[:-1], [True] )))
    return [data_mat[i:j] for i,j in zip(idx[:-1], idx[1:])]

Тайминги -

В примере у нас была подматрица максимальной длины 6. Итак, давайте расширимся до более крупного случая, сохранив его по той же схеме -

In [442]: a = np.random.randint(0,100000,(6*100000,4)); a[:,0].sort()

In [443]: %timeit split_app(a)
10 loops, best of 3: 88.8 ms per loop

In [444]: %timeit zip_app(a)
10 loops, best of 3: 40.2 ms per loop

In [445]: a = np.random.randint(0,1000000,(6*1000000,4)); a[:,0].sort()

In [446]: %timeit split_app(a)
1 loop, best of 3: 917 ms per loop

In [447]: %timeit zip_app(a)
1 loop, best of 3: 414 ms per loop

Вы можете сделать это с помощью логического индексирования.

unique_ids = np.unique(data_mat[:, 0])
masks = np.equal.outer(unique_ids, data_mat[:, 0])
for mask in masks:
    tmp_mat = data_mat[mask]
    # do something with tmp_mat ...
    print(tmp_mat)

Если уникальные идентификаторы уже сгруппированы, вы можете сделать это с помощью np.split, аналогично @Divakar.

idx = np.flatnonzero(data_mat[1:, 0] != data_mat[:-1, 0]) + 1
for tmp_mat in np.split(data_mat, idx):
    # do something with tmp_mat ...
    print(tmp_mat)