zip_longest для массивов NumPy

Question

Я пытаюсь объединить массивы numpy разной длины, как можно было бы эквивалентно использовать списки с itertools.zip_longest.Скажем, у меня есть:

a = np.array([1, 5, 9, 13])
b = np.array([2, 6])

С itertools можно чередовать эти два массива, используя chain и zip_longest, и заполнить пропущенные значения с помощьюскажем 0:

from itertools import chain, zip_longest
list(chain(*zip_longest(*[a, b], fillvalue=0)))
# [1, 2, 5, 6, 9, 0, 13, 0]

Есть ли простой способ сделать это с помощью numpy, который мне не хватает?

Answer 1

Вот почти векторизованный -

# https://stackoverflow.com/a/38619350/3293881 @Divakar
def boolean_indexing(v):
    lens = np.array([len(item) for item in v])
    mask = lens[:,None] > np.arange(lens.max())
    out_dtype = np.result_type(*[arr.dtype for arr in v])
    out = np.zeros(mask.shape,dtype=out_dtype)
    out[mask] = np.concatenate(v)
    return out

v = [a,b] # list of all input arrays
out = boolean_indexing(v).ravel('F')

Пробный прогон -

In [23]: a = np.array([1, 5, 9, 13])
    ...: b = np.array([2, 6])
    ...: c = np.array([7, 8, 10])
    ...: v = [a,b,c]

In [24]: boolean_indexing(v).ravel('F')
Out[24]: array([ 1,  2,  7,  5,  6,  8,  9,  0, 10, 13,  0,  0])

Answer 2

Я думаю, что сделал бы это так:

import numpy as np

def chain_zip_longest(*arrs, fillvalue=0, dtype=None):
    arrs = [np.asarray(arr) for arr in arrs]
    if not arrs:
        return np.array([])
    n = len(arrs)
    dtype = dtype or np.find_common_type([arr.dtype for arr in arrs], [])
    out = np.full(n * max(len(arr) for arr in arrs), fillvalue, dtype=dtype)
    for i, arr in enumerate(arrs):
        out[i:i + n * len(arr):len(arrs)] = arr
    return out

print(chain_zip_longest([1, 2], [3, 4, 5], [6]))
# [1 3 6 2 4 0 0 5 0]