Repeat Numpy Array Variable Number of Times Depending on Row

user_8206432 0
Viewed 301

I am trying to create a numpy array using a pandas data frame. One column, named "repeat", in the data frame tells what value to repeat and the other column, named "times", tells how many times that value should be repeated.

Here is what I have tried:

import pandas as pd
import numpy as np

df = pd.DataFrame({'repeat': [1, 4, 3, 2], 'times': [2, 5, 4, 1]})
np.repeat(df['repeat'].values, df['times'])

However, this outputs the following:

array([1, 1, 4, 4, 4, 4, 4, 3, 3, 3, 3, 2])

This is what I want:

desired = np.array([[1, 1], [4, 4, 4, 4, 4], [3, 3, 3, 3], [2]])

array([list([1, 1]), list([4, 4, 4, 4, 4]), list([3, 3, 3, 3]), list([2])], dtype=object)

How can I repeat each value the corresponding number of times along the row efficiently? Is there a good numpy or pandas solution here?

pandas repeat numpy arrays
3 Answers

You can get an array of arrays by splitting the original result at indices based on the size:

r = [1, 4, 3, 2]
t = [2, 5, 4, 1]
np.split(np.repeat(r, t), np.cumsum(t[:-1]))

This will return a list of arrays, which is arguably a much more efficent way to store the data than an array of lists. Remember, arrays are good for contiguous blocks of numerical data, while lists are good for arbitrary objects.

user_29887300

Let's try repeating the index and groupby:

(df.loc[df.index.repeat(df['times']),'repeat']
   .groupby(level=0)
   .agg(list).values
)

Output

array([list([1, 1]), list([4, 4, 4, 4, 4]), list([3, 3, 3, 3]), list([2])],
      dtype=object)
user_42384080

I will do

import itertools
[list(itertools.repeat(x, y)) for x, y  in zip(df['repeat'], df['times'])]
Out[301]: [[1, 1], [4, 4, 4, 4, 4], [3, 3, 3, 3], [2]]

Some timing

#df=pd.concat([df]*100)

%timeit [list(itertools.repeat(x, y)) for x, y  in zip(df['repeat'], df['times'])]
178 µs ± 521 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit df.loc[df.index.repeat(df['times']),'repeat'].groupby(level=0).agg(list).values
20.2 ms ± 562 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

def func():
    r = df['repeat']
    t = df['times']
    return np.split(np.repeat(r, t), np.cumsum(t[:-1]))
%timeit func()
45.1 ms ± 846 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
user_79645270
Related