splinter - 9 months ago 43

Python 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`

`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))`

`DataFrame`

Answer Source

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]
```