Resize a numpy array of lists so that so that the lists all have the same length and dtype of numpy array can be inferred correctly

I currently have the following dataframe

data = {'col_a': [['a', 'b'], ['a', 'b', 'c'], ['a'], ['a', 'b', 'c', 'd'], ['a', 'b', 'c'], ['a', 'b', 'c', 'd']],
        'col_b':[[1, 3], [1, 0, 0], [4], [1, 1, 2, 0], [0, 0, 5], [3, 1, 2, 5]]}
df= pd.DataFrame(data)

Suppose I work with col_a, I want to resize the lists in col_a in a vectorized manner so that the length of all the sub lists = max length of largest list and I want to fill the empty values with 'None' in the case of col_a. I want the final output to look as follows

                   col_a               col_b
0     [a, b, None, None]    [1, 3, nan, nan]
1        [a, b, c, None]      [1, 0, 0, nan]
2  [a, None, None, None]  [4, nan, nan, nan]
3           [a, b, c, d]        [1, 1, 2, 0]
4        [a, b, c, None]      [0, 0, 5, nan]
5           [a, b, c, d]        [3, 1, 2, 5]

So far I have done the following

# Convert the column to a NumPy array with object dtype
col_np = df['col_a'].to_numpy()

# Find the maximum length of the lists using NumPy operations
max_length = np.max(np.frompyfunc(len, 1, 1)(col_np))

# Create a mask for padding
mask = np.arange(max_length) < np.frompyfunc(len, 1, 1)(col_np)[:, None]

# Pad the lists with None where necessary
result = np.where(mask, col_np, 'None')

This results in the following error
ValueError: operands could not be broadcast together with shapes (6,4) (6,) ()

I feel like I'm close but there's something that I'm missing here. Please note that only vectorized solutions will be marked as the answer.

Only vectorized solutions will be marked as the answer. -> that's too bad because no (true) vectorized approach is possible with an array of lists. To this extent, np.frompyfunc is certainly not truly vectorized.

If by "vectorized" you mean without explicit python loop, you could use:

df['out_a'] = pd.Series(pd.DataFrame(df['col_a'].to_numpy().tolist()).to_numpy().tolist())

An alternative with an explicit loop would be:

size = df['col_a'].str.len().max()

df['out_a'] = [l+[None]*(size-len(l)) for l in df['col_a']]

Output:

          col_a         col_b                  out_a
0        [a, b]        [1, 3]     [a, b, None, None]
1     [a, b, c]     [1, 0, 0]        [a, b, c, None]
2           [a]           [4]  [a, None, None, None]
3  [a, b, c, d]  [1, 1, 2, 0]           [a, b, c, d]
4     [a, b, c]     [0, 0, 5]        [a, b, c, None]
5  [a, b, c, d]  [3, 1, 2, 5]           [a, b, c, d]

timings

For small lists the "vectorized" and loop solution have very similar timings.

Here for lists with 1 to 10 items:

However, when the size of lists increases, the python loop become more efficient.

For list with 0 to 50 items:

0 to 200 items:

0 to 2000 items:

Code used for the timings:

import pandas as pd
import perfplot
import numpy as np

def pandas_vectorized(df):
    df['out_a'] = pd.Series(pd.DataFrame(df['col_a'].to_numpy().tolist()).to_numpy().tolist())
    
def python_loop(df):
    size = df['col_a'].str.len().max()
    df['out_a'] = [l+[None]*(size-len(l)) for l in df['col_a']]

MAX_LIST_SIZE = 2000
    
perfplot.show(
    setup=lambda n: pd.DataFrame({'col_a': [['x']*n for n in np.random.randint(0, MAX_LIST_SIZE, size=n)]}),
    kernels=[pandas_vectorized, python_loop],
    n_range=[2**k for k in range(1, 18)],  # upper bound was 22 for small lists
    xlabel="len(df)",
    equality_check=None,
    max_time=10,
)

This is basically a padding question - expanding lists to the matching length by adding fill values. This has come up a number times.

itertools has a variant on zip that is useful for this:

In [273]: from itertools import zip_longest

Your column:

In [274]: x=df.col_a.to_numpy()
In [275]: x
Out[275]: 
array([list(['a', 'b']), list(['a', 'b', 'c']), list(['a']),
       list(['a', 'b', 'c', 'd']), list(['a', 'b', 'c']),
       list(['a', 'b', 'c', 'd'])], dtype=object)

This pads, but it's the transpose of what we want:

In [276]: list(zip_longest(*x))
Out[276]: 
[('a', 'a', 'a', 'a', 'a', 'a'),
 ('b', 'b', None, 'b', 'b', 'b'),
 (None, 'c', None, 'c', 'c', 'c'),
 (None, None, None, 'd', None, 'd')]

We could use an array transpose, but there's a list idiom that does the same:

In [277]: list(zip(*(zip_longest(*x))))
Out[277]: 
[('a', 'b', None, None),
 ('a', 'b', 'c', None),
 ('a', None, None, None),
 ('a', 'b', 'c', 'd'),
 ('a', 'b', 'c', None),
 ('a', 'b', 'c', 'd')]

And the same thing applied to the numeric column, with nan fill:

In [281]: list(zip(*(zip_longest(*df.col_b, fillvalue=np.nan))))
Out[281]: 
[(1, 3, nan, nan),
 (1, 0, 0, nan),
 (4, nan, nan, nan),
 (1, 1, 2, 0),
 (0, 0, 5, nan),
 (3, 1, 2, 5)]

For large lists of lists there are faster padding methods, but this zip_longest is one of the easier ones to remember and use.

edit

With your code, mask is a good looking array, telling where we want the fill values:

In [285]: mask
Out[285]: 
array([[ True,  True, False, False],
       [ True,  True,  True, False],
       [ True, False, False, False],
       [ True,  True,  True,  True],
       [ True,  True,  True, False],
       [ True,  True,  True,  True]])

It's (6,4). To be used in where where have to expand col_np to (6,1):

In [286]: result = np.where(mask, col_np[:,None], 'None')
     ...: 

But the result is hardly what we want:

In [287]: result
Out[287]: 
array([[list(['a', 'b']), list(['a', 'b']), 'None', 'None'],
       [list(['a', 'b', 'c']), list(['a', 'b', 'c']),
        list(['a', 'b', 'c']), 'None'],
       [list(['a']), 'None', 'None', 'None'],
       [list(['a', 'b', 'c', 'd']), list(['a', 'b', 'c', 'd']),
        list(['a', 'b', 'c', 'd']), list(['a', 'b', 'c', 'd'])],
       [list(['a', 'b', 'c']), list(['a', 'b', 'c']),
        list(['a', 'b', 'c']), 'None'],
       [list(['a', 'b', 'c', 'd']), list(['a', 'b', 'c', 'd']),
        list(['a', 'b', 'c', 'd']), list(['a', 'b', 'c', 'd'])]],
      dtype=object)

However there is a way of using that mask to fill an array. First make the target with the right shape and dtype:

In [288]: result = np.empty((6,4),object)
In [289]: result
Out[289]: 
array([[None, None, None, None],
       [None, None, None, None],
       [None, None, None, None],
       [None, None, None, None],
       [None, None, None, None],
       [None, None, None, None]], dtype=object)
In [290]: result[mask] = np.hstack(col_np)
In [291]: result
Out[291]: 
array([['a', 'b', None, None],
       ['a', 'b', 'c', None],
       ['a', None, None, None],
       ['a', 'b', 'c', 'd'],
       ['a', 'b', 'c', None],
       ['a', 'b', 'c', 'd']], dtype=object)

This fill works because result[mask] is a flat array of values:

In [292]: result[mask]
Out[292]: 
array(['a', 'b', 'a', 'b', 'c', 'a', 'a', 'b', 'c', 'd', 'a', 'b', 'c',
       'a', 'b', 'c', 'd'], dtype=object)

Overall this approach is quite close to one of the more clever, and faster, padding methods that have been proposed. I believe Warren Weckesser was the original source.

And the numeric column

In [300]: result = np.full((6,4),np.nan)
In [301]: result
Out[301]: 
array([[nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan],
       [nan, nan, nan, nan]])
In [302]: result[mask] = np.hstack(df.col_b)
In [303]: result
Out[303]: 
array([[ 1.,  3., nan, nan],
       [ 1.,  0.,  0., nan],
       [ 4., nan, nan, nan],
       [ 1.,  1.,  2.,  0.],
       [ 0.,  0.,  5., nan],
       [ 3.,  1.,  2.,  5.]])

The mask can be created with only one list comprehension - to get the length of each list.

In [318]: lens = np.array([len(i) for i in x])
In [319]: np.max(lens)
Out[319]: 4
In [320]: mask1 = np.arange(np.max(lens))<lens[:,None]
In [321]: mask1
Out[321]: 
array([[ True,  True, False, False],
       [ True,  True,  True, False],
       [ True, False, False, False],
       [ True,  True,  True,  True],
       [ True,  True,  True, False],
       [ True,  True,  True,  True]])

timings

With Mozway's dataframe expansion:

In [325]: timeit pd.DataFrame(df['col_a'].to_numpy().tolist())
296 µs ± 3.17 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
In [326]: timeit pd.Series(pd.DataFrame(df['col_a'].to_numpy().tolist()).to_numpy().toli
     ...: st())
567 µs ± 52.8 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Packaging the mask approach into a function:

def foo(df):
    x = df.col_a.to_numpy()
    lens = np.array([len(i) for i in x])
    mask = np.arange(np.max(lens))<lens[:,None]
    result = np.empty(mask.shape, dtype=object)
    result[mask] = np.hstack(x)
    return result

Time to get the padded array:

In [337]: timeit foo(df)
74.6 µs ± 3.45 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

and converting that array to a Series:

In [341]: timeit pd.Series(foo(df).tolist())
299 µs ± 68.6 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

and the zip_longest approach (fastest for this size)

In [344]: pd.Series(list(zip(*(zip_longest(*df.col_a)))))
Out[344]: 
0       (a, b, None, None)
1          (a, b, c, None)
2    (a, None, None, None)
3             (a, b, c, d)
4          (a, b, c, None)
5             (a, b, c, d)
dtype: object
In [345]: timeit list(zip(*(zip_longest(*df.col_a))))
19.7 µs ± 151 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
In [346]: timeit pd.Series(list(zip(*(zip_longest(*df.col_a)))))
165 µs ± 29.6 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)