zip_longest for numpy arrays

Asked by At

I'm trying to combine different length numpy arrays as one could equivalently do using lists with itertools.zip_longest. Say I have:

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

With itertools one could interleave these two arrays using chain and zip_longest, and fill the missing values with say 0:

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

Is there a simple way to do this using numpy that I'm missing?

2 Answers

1
Divakar On Best Solutions

Here's an almost vectorized one -

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

Sample run -

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])
1
jdehesa On

I think I'd do that like this:

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]