Creating a new column consisting of lists in a DataFrame using pandas

Question

Given the following DataFrame:

   t
0  3
1  5

I would like to create a new column where wach entry is a list which is a function of the row it is in. In particular it should have a list with all positive integers which not greater than the entry in column t. So the output should be:

   t  newCol
0  3  [1,2,3]
1  5  [1,2,3,4,5]

In other words, I want to apply list(range(1,t+1)) to each row. I know how to do it in a loop, but have a long DataFrame, so I am looking for speed. Thank you.

df['newCol'] = df.t.map(np.arange) + 1

MaxU - stand with Ukraine
– MaxU - stand with Ukraine

2016-11-13 23:36:42 +00:00
Commented Nov 13, 2016 at 23:36 — MaxU - stand with Ukraine
– MaxU - stand with Ukraine, Commented Nov 13, 2016 at 23:36

Divakar · Accepted Answer · 2016-11-13 23:45:46Z

2

Here's a vectorized approach using NumPy methods -

a = df.t.values
idx = a.cumsum()
id_arr = np.ones(idx[-1],dtype=int)
id_arr[idx[:-1]] = -a[:-1]+1
df['newCol'] = np.split(id_arr.cumsum(),idx[:-1])

Sample run -

In [76]: df
Out[76]: 
   t                 newCol
0  4           [1, 2, 3, 4]
1  3              [1, 2, 3]
2  7  [1, 2, 3, 4, 5, 6, 7]
3  2                 [1, 2]
4  5        [1, 2, 3, 4, 5]
5  3              [1, 2, 3]

answered Nov 13, 2016 at 23:45

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

piRSquared · Accepted Answer · 2016-11-14 02:14:24Z

this is so very close to @Divakar's answer but I believe a tad bit more intuitive.

get values for quicker numpy access
v = df.t.values
[3 5]
get cumulative sums of v
cumsum = v.cumsum()
[3 8]
get some differences
used to track splits and take differences later
diffs = cumsum - v
[0 3]
compile a big cumulative sum
This is the starting point for the final values
prevals = np.ones(cumsum[-1], dtype=int).cumsum()
[1 2 3 4 5 6 7 8]
finally, split and put
df['new_col'] = np.split(prevals - np.repeat(diffs, v), diffs[1:])

all together

df = pd.DataFrame(dict(t=[4, 3, 7, 2, 5, 3]))
v = df.t.values
cumsum = v.cumsum()
diffs = cumsum - v
prevals = np.ones(cumsum[-1], dtype=int).cumsum()
df['new_col'] = np.split(prevals - np.repeat(diffs, v), diffs[1:])

Collectives™ on Stack Overflow

Creating a new column consisting of lists in a DataFrame using pandas

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related