Разделение массива NumPy на несколько массивов на каждое количество ячеек

Question

У меня есть массив массивов 1D, который я хочу разделить на несколько массивов в зависимости от количества ячеек (x) в каждом массиве.

Например, если это массив:

[0, 2, 7, 9, 13]

и x=2

Я хочу получить в результате следующие массивы:

[0,2]
[7,9]
[13]

Мне кажется, что есть простой способ сделать это без использования циклов,но я не смог его найти.

Спасибо!

Answer 1

import numpy as np
A = np.asarray([0, 2, 7, 9, 13])    
x = 2
np.split(A, range(x, len(A), x))  
# [array([0, 2]), array([7, 9]), array([13])]

Обратите внимание, что, как @Divakar упомянул в своем комментарии, вам лучше справляться с производительностью, используя циклы (или, в данном случае, понимание списка).

def f(seq, n):
    return [seq[i:min(i + n, len(seq))] for i in range(0, len(seq), n)]

def g(seq, n):
    return np.split(seq, range(n, len(A), n))

%timeit f(range(12345), 109)
65.2µs ± 468ns (mean +- std. dev. of 7 runs, 10000 loops each)

%timeit g(range(12345), 109)
1.32ms ± 5.21 µs (mean +- std. dev. of 7 runs, 10000 loops each)

Редактировать

Поскольку разрезание последнего индекса естественным образом будет срезаться до последнего элемента, мы можем отменить проверку

def f(seq, n):
    return [seq[i : i + n] for i in range(0, len(seq), n)]