python - 如何根据用户在运行时选择的逻辑创建新的 pandas 列

Question

我希望能够根据用户选择的公式在数据框中创建一个新列。当在运行之前知道公式选择时，答案是显而易见的，但我不确定在直到运行时才知道公式的情况下如何进行。

数据框中的列数可能会有所不同，每个公式的变量/操作数也可能会有所不同，因此 lambda 函数似乎不合适

---

前任。df在“A”列中具有整数值，在“B”列中具有浮点值。有3个公式，

• 1- 将 A 中的值添加到 B 中的值
• 2- 从 A 中的值的平方中减去 B 中的值
• 3- 将 A 中的值添加到 B 中的值

用户输入1，2或3在运行时，相应的公式用于创建具有适当值的列“C”。

标签： pythonpandasdynamic

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 中作为对象的函数。