Как создать новый столбец панд на основе логики, выбранной во время выполнения пользователем

Question

Я хочу иметь возможность создавать новый столбец в кадре данных на основе выбранной пользователем формулы. Ответ очевиден, когда выбор формулы известен до запуска, но я не уверен, как поступить в случае, когда формула неизвестна до времени выполнения.

Количество столбцов в кадре данных может варьироваться, а также может варьироваться число переменных / операций в формуле, поэтому лямбда-функция кажется неуместной

---

Пример. df имеет целочисленные значения в столбце «A» и значения с плавающей запятой в столбце «B». Есть 3 формулы:

• 1 - добавить значение в A к значению в B
• 2 - вычесть значение в B из квадрата значения вA
• 3 - добавить значение в A к значению в B

, которое пользователь вводит 1, 2 или 3 во время выполнения, и соответствующееформула используется для создания столбца «С» с соответствующими значениями.

Answer 1

def get_formula(user_input):
  formula = None
  if user_input == 1:
    formula = lambda x: x['A'] + x['B']
  elif user_input == 2:
    formula = lambda x: x['A']**2 - x['B']
  elif user_input == 3:
    # Your conditions for user inputs 1 or 3 seem to be the same.
    formula = lambda x: x['A'] + x['B']
  else:
    # Error out
    pass
  return formula

df = pd.DataFrame({'A': list(range(5)), 'B': list(range(5))})

# Get user input (assumes it's an integer)
user_input = int(input('Enter formula #: '))
# Get formula based on input
fn = get_formula(user_input)
# Assign new column 'C' based on formula
df.assign(C=fn)

# Assuming user input is 1, this outputs:
   A  B  C
0  0  0  0
1  1  1  2
2  2  2  4
3  3  3  6
4  4  4  8

Вы можете использовать DataFrame.assign вместе с формулой, определенной во время выполнения. Это позволяет использовать функции, являющиеся объектами в Python.