如何基于days_to_last_follow_up值条件生成survival行值？

需求说明 #

当days_to_last_follow_up的行值大于等于1825时，将survival行值赋值为0；若该值小于1825或为NA，则将survival行值赋值为1。

现有代码问题 #

当前实现逻辑错误：函数返回的是布尔比较结果而非目标赋值，且未处理NA场景，无法得到预期输出：

# Long-term survival >= 5 years (1825 days)
# Short-term survival < 5 years OR NA
def survival_status(col):
  if col.loc["days_to_last_follow_up"] >= 1825:
    return col.loc["survival"] == 0 # lts
  else:
    return col.loc["survival"] == 1 # non-lts

clinical.loc["survival"] = clinical.apply(survival_status, axis=0)

输入数据（clinical.iloc[0:4,0:4]） #

import pandas as pd

pd.DataFrame({'TCGA-2K-A9WE-01': {'admin.batch_number': '398.45.0',
  'age': '53',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '207.0'},
 'TCGA-2Z-A9J1-01': {'admin.batch_number': '398.45.0',
  'age': '71',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '2298.0'},
 'TCGA-2Z-A9J3-01': {'admin.batch_number': '398.45.0',
  'age': '67',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': float('nan')},
 'TCGA-2Z-A9J6-01': {'admin.batch_number': '398.45.0',
  'age': '60',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '1731.0'}})

预期输出 #

pd.DataFrame({'TCGA-2K-A9WE-01': {'admin.batch_number': '398.45.0',
  'age': '53',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '207.0',
  'survival': '1'},
 'TCGA-2Z-A9J1-01': {'admin.batch_number': '398.45.0',
  'age': '71',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '2298.0',
  'survival': '0'},
 'TCGA-2Z-A9J3-01': {'admin.batch_number': '398.45.0',
  'age': '67',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': float('nan'),
  'survival': '1'},
 'TCGA-2Z-A9J6-01': {'admin.batch_number': '398.45.0',
  'age': '60',
  'days_to_initial_pathologic_diagnosis': '0',
  'days_to_last_follow_up': '1731.0',
  'survival': '1'}})
# 注：原预期输出中TCGA-2Z-A9J6-01的survival应为1，因1731<1825符合赋值1的条件

正确实现方案 #

方案1：修正apply函数实现 #

# 先将天数列转为数值类型，避免字符串比较错误
clinical.loc["days_to_last_follow_up"] = pd.to_numeric(clinical.loc["days_to_last_follow_up"], errors='coerce')

def survival_status(col):
    days = col.loc["days_to_last_follow_up"]
    # NA或小于1825返回1，否则返回0
    if pd.isna(days) or days < 1825:
        return '1'
    else:
        return '0'

clinical.loc["survival"] = clinical.apply(survival_status, axis=0)

方案2：向量化操作（更高效，推荐） #

# 转换数据类型
clinical.loc["days_to_last_follow_up"] = pd.to_numeric(clinical.loc["days_to_last_follow_up"], errors='coerce')

# 用numpy.where实现批量条件赋值
import numpy as np
clinical.loc["survival"] = np.where(
    (clinical.loc["days_to_last_follow_up"] >= 1825) & (~pd.isna(clinical.loc["days_to_last_follow_up"])),
    '0',
    '1'
)

关键修正点 #

1. 数据类型转换：输入中days_to_last_follow_up为字符串类型，必须转为数值类型才能正确进行大小比较。
2. NA值处理：新增NA判断逻辑，确保NA场景下赋值为1。
3. 直接返回目标值：函数不再返回布尔比较表达式，而是直接返回'0'或'1'，匹配预期输出的字符串格式。

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