Applying a subscripting function within a data frame-CodePudding

I have a a column with lists of numbers and have written a function to take the even numbers of the list if the first number of the list is even and the odd numbers if the first number of the list is odd. However, I am running into a TypeError: 'NoneType' object is not subscriptable when running my script.

def remove(val):
  even = []
  odd = []
  if val[0] % 2 == 0:
    for i in val:
      if (i % 2 == 0):
        even.append(i)
    return even
  else:
    for i in val:
      if (i % 2 != 0):
        odd.append(i)
    return odd

df.B = df.A.apply(remove)

The expected result would be something like below:

       A               B    
[1,2,3,4,5,6]       [1,3,5]
[2,3,4,5,6,7]       [2,4,6]
      ...             ...

I am realizing that there is most likely a more eloquent way to build my function with list comprehension but this seemed the most obvious.

CodePudding user response：

You are not using .apply() correctly. try:

X = pd.DataFrame({
        'A':[[1,2,3,4,5,6],[2,3,4,5,6,7]]})
X['B'] = X.apply(lambda x: remove(x[0]),axis=1)

Also, if you'd like a solution which doesn't use an extrernal function:

X = pd.DataFrame({
        'A':[[1,2,3,4,5,6],[2,3,4,5,6,7]]})

X['B'] = X.apply(lambda x: [i for i in x[0] if i % 2 == 0] if x[0][0]%2 == 0 else [i for i in x[0] if i % 2 == 1],axis=1)