Numpy, как разбить список на несколько частей

Question

Я пытаюсь разбить массив с кусочками на куски с фиксированным размером и заполнить последний с 0. Например: [1,2,3,4,5,6,7] на куски 3 возвращает [[1,2,3],[4,5,6],[7,0,0]].

ФункцияЯ написал:

def makechunk(lst, chunk):
    result = []
    for i in np.arange(0, len(lst), chunk):
        temp = lst[i:i + chunk]
        if len(temp) < chunk:
            temp = np.pad(temp, (0, chunk - len(temp)), 'constant')
        result.append(temp)
    return result

Это работает, но при работе с массивом большого размера это очень медленно.Что такое более numpy-ish и векторизованный способ сделать это?

Answer 1

Сравнение времени @ решения Седрика Пуле (все ему благодарны, см. Его ответ) (с добавленным разбиением массива, чтобы он возвращал результат по желанию) с другим подходом numpy, о котором я думал вначале (создать массив нулей и вставить данные на месте):

import time

import numpy as np

def time_measure(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        stop = time.time()
        print(f"Elapsed time: {stop-start}")
        return result

    return wrapper


@time_measure
def pad_and_chunk(array, chunk_size: int):
    padded_array = np.zeros(len(array) + (chunk_size - len(array) % chunk_size))
    padded_array[: len(array)] = array
    return np.split(padded_array, len(padded_array) / chunk_size)


@time_measure
def resize(array, chunk_size: int):
    array.resize(len(array) + (chunk_size - len(array) % chunk_size), refcheck=False)
    return np.split(array, len(array) / chunk_size)

@time_measure
def makechunk4(l, chunk):
    l.resize((math.ceil(l.shape[0] / chunk), chunk), refcheck=False)
    return l.reshape(chunk, -1)


if __name__ == "__main__":
    array = np.random.rand(1_000_000)

    ret = pad_and_chunk(array, 3)
    ret = resize(array, 3)
    ret = makechunk4(array, 3)

EDIT-EDIT

Собирая все возможные ответы, действительно, np.split ужасно медленный по сравнению с изменением формы.

Elapsed time: 0.3276541233062744
Elapsed time: 0.3169224262237549
Elapsed time: 1.8835067749023438e-05

Способ заполнения данными не важен, это разделение, занимающее большую часть времени.

Answer 2

Использование функции resize () должно делать то, что вам нужно:

l = np.array([1,2,3,4,5,6,7])
l.resize((3,3), refcheck=False)

(Редактировать: mea culpa, проблема понедельника с переназначением)

@ J: Изменение размера увеличивает скорость примерно на5 раз для np.arange (0,44100) на куски по 512.

import math
def makechunk4(lst, chunk):
    l = lst.copy()
    l.resize((math.ceil(l.shape[0]/chunk),chunk), refcheck=False)
    return l

Answer 3

Решение с использованием numpy

Я предполагаю, что размер порции равен 3, и создал случайный вход массива длиной 10 в x.

# Chunk size
chunk = 3
# Create array
x = np.arange(10)

Сначала убедитесь, что массив заполнен нулями. Далее вы можете использовать reshape для создания массива массивов.

# Pad array
x = np.pad(x, (0, chunk - (x.shape[0]%chunk)), 'constant')
# Divide into chunks
x = x.reshape(-1, chunk)

При желании вы можете получить массив numpy в виде списка

x = x.tolist()

Answer 4

в рецептах itertools есть рецепт для grouper:

from itertools import zip_longest
import numpy as np

array = np.array([1,2,3,4,5,6,7])

def grouper(iterable, n, fillvalue=None):
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

res = list(grouper(array, 3, fillvalue=0))
# [(1, 2, 3), (4, 5, 6), (7, 0, 0)]

, если вам нужно, чтобы список составлял list с, а не tuple с:

def grouper(iterable, n, fillvalue=None):
    args = [iter(iterable)] * n
    return (list(item) for item in zip_longest(*args, fillvalue=fillvalue))