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

我想我会这样做：

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]

Answer 2

这几乎是矢量化的-

# 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])