Remove n elements from start of a list in pandas column, where n is the value in another column

user_8868620 0
Viewed 81

Say I have the following DataFrame:

a = [[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15]]
b = [3,1,2]
df = pd.DataFrame(zip(a,b), columns = ['a', 'b'])

df:

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

How can I remove the first n elements from each list in column a, where n is the value in column b.

The result I would expect for the above df is:

                      a  b
0                [4, 5]  3
1         [7, 8, 9, 10]  1
2          [13, 14, 15]  2

I imagine the answer revolves around using .apply() and a lambda function, but I cannot get my head around this one!

pandas dataframe python
2 Answers

Try:

df["a"] = df.apply(lambda x: x["a"][x["b"] :], axis=1)
print(df)

Prints:

               a  b
0         [4, 5]  3
1  [7, 8, 9, 10]  1
2   [13, 14, 15]  2
user_100359850

Try this:

df['a'] = df.apply(lambda row: row['a'][row['b']:], axis=1)

Output:

               a    b
0          [4, 5]   3
1   [7, 8, 9, 10]   1
2    [13, 14, 15]   2
user_17405770
Related