How to vectorize dynamically sized numpy arrays in pandas

user_8206432 0
Viewed 134

Currently, I am using the apply method on my dataframe in order to create a computed column which contains lists of variable sizes (depending on the value in the length column).

Is there a way to create a column with variable-sized lists more efficiently with pandas?

import pandas as pd
import numpy as np

df = pd.DataFrame({'a': [1, 2, 3, 4, 5], 'b': [
                  6, 7, 8, 9, 0], 'length': [3, 5, 7, 9, 3]})
df['computed'] = df.apply(
    lambda x: np.array([x['a'], x['b']] + [x['b'] + i for i in range(1, x['length'] - 1)]), axis=1)

Desired output (works with code above, but slow):

   a  b  length                            computed
0  1  6       3                           [1, 6, 7]
1  2  7       5                    [2, 7, 8, 9, 10]
2  3  8       7           [3, 8, 9, 10, 11, 12, 13]
3  4  9       9  [4, 9, 10, 11, 12, 13, 14, 15, 16]
4  5  0       3                           [5, 0, 1]
pandas apply dataframe numpy python
2 Answers

try this,

df['computed']= [[a]+[b]+list(np.arange(b+1, length)) for a, b, length in zip(df.a, df.b, (df.b) + df.length-1)]

o/P:

   a  b  length                            computed
0  1  6       3                           [1, 6, 7]
1  2  7       5                    [2, 7, 8, 9, 10]
2  3  8       7           [3, 8, 9, 10, 11, 12, 13]
3  4  9       9  [4, 9, 10, 11, 12, 13, 14, 15, 16]
4  5  0       3                           [5, 0, 1]
user_46848610

Not sure if this is what you were looking for but, if it's too slow you can always try multiprocessing:

import pandas as pd
import numpy as np
from multiprocessing import Pool

def parallelize(df, func, n_cores=4):
    df_split = np.array_split(df, n_cores)
    pool = Pool(n_cores)
    df = pd.concat(pool.map(func, df_split))
    pool.close()
    pool.join()
    return df

def func(df):
    df['computed'] = df.apply(lambda x: np.array([x['a'], x['b']] + [x['b'] + i for i in range(1, x['length'] - 1)]), axis=1)
    return df

df = pd.DataFrame({'a': [1, 2, 3, 4, 5], 'b': [6, 7, 8, 9, 0], 'length': [3, 5, 7, 9, 3]})
df = parallelize(df, func)

(for small values of length it will be less efficient than your original code)

user_47577150
Related