Pandas根据Year列匹配对应人口列创建DataFrame新列

这是我首次在Stack Overflow提问，若格式存在瑕疵还请谅解。
我在完成多个DataFrame拼接后，始终无法实现基于其他列（df["2013 Pop"]、df["2014 Pop"]等）的行值新建df["population"]列的需求：当事件发生在对应年份时，需要提取相同年份Pop列的对应人口值，填充到新的df["population"]列中。比如df["Year"] == 2014时，取df["2014 Pop"]的同行值填入新列。

我的DataFrame结构示例如下：

d = {
  "Year" : [2013,2014,2015...],
  "State" : ["Louisana", "Texas", "California"... ],
  "City" : ["New Orleans", "Dallas", "Sacramento"...],
  "Number Killed" : [4,6,2,4],
  "Safety Grade" : ["A", "B", "C", "D"...],
  "2013 Pop" : [421329, 232321, 2454543....],
  "2014 Pop" : [454545, 655654, 3421342....],
  "2015 Pop" : [142314, 454355, 4324323....],
  "Incident Date(datetime dtype)" : ["12-29-2014", "3-12-2017"...]
}
df = pd.DataFrame(d)

我尝试过mapping、loc、apply等方法均未成功，原本认为带条件判断的自定义函数思路可行，但运行时持续报错。我编写的函数和调用代码如下：

def categorise(row):
  if row["Year"] == 2014:
    return df["2014 Pop"]
  elif row["Year"] == 2015:
    return df["2015 Pop"]
  elif row["Year"] == 2016:
    return df["2016 Pop"]
  elif row["Year"] == 2017:
    return df["2017 Pop"]
  else:
    return "NONE"

df["Population"] = df.apply(lambda row : categorise(row), axis = 1)

执行后抛出错误：ValueError: Wrong number of items passed 3609 (length of the df), placement implies 1

核心问题出在自定义函数的返回值：apply(axis=1)是逐行遍历操作，每次传入的row是当前行的Series对象，你返回的是整个DataFrame的整列数据（比如df["2014 Pop"]长度为3609，等于整个表的行数），而非当前行对应的单个数值，pandas无法将长度为3609的序列填入单个单元格，因此抛出长度不匹配错误。

方案1：numpy.select 向量化实现（优先推荐） #

使用numpy的条件选择函数做向量化匹配，不需要逐行遍历，运行效率比apply高数十倍，适合大体量数据：

import numpy as np

# 按顺序写匹配条件、对应返回的列
conditions = [
    df["Year"] == 2013,
    df["Year"] == 2014,
    df["Year"] == 2015,
    df["Year"] == 2016,
    df["Year"] == 2017
]
pop_cols = [
    df["2013 Pop"],
    df["2014 Pop"],
    df["2015 Pop"],
    df["2016 Pop"],
    df["2017 Pop"]
]

# 匹配到对应条件就取对应列的同行值，未匹配到则返回"NONE"
df["Population"] = np.select(conditions, pop_cols, default="NONE")

方案2：修正原有apply自定义函数 #

如果习惯使用apply写法，只需要把函数中返回整列的写法改为返回当前行的对应值即可，同时可以去掉多余的lambda嵌套：

def categorise(row):
  if row["Year"] == 2014:
    return row["2014 Pop"] # 注意这里取当前row的对应值，不是全局df的整列
  elif row["Year"] == 2015:
    return row["2015 Pop"]
  elif row["Year"] == 2016:
    return row["2016 Pop"]
  elif row["Year"] == 2017:
    return row["2017 Pop"]
  else:
    return "NONE"

# 直接传入函数对象即可，不需要套lambda
df["Population"] = df.apply(categorise, axis=1)

方案3：动态列名匹配（适合年份较多的场景） #

如果涉及的年份范围很大，逐个写if判断代码冗余度高，可以直接通过行内的Year值拼接列名，动态取对应值：

df["Population"] = df.apply(
    lambda row: row.get(f"{int(row['Year'])} Pop", "NONE"),
    axis=1
)

内容的提问来源于stack exchange，提问作者Nickguild1993