Python中获取Cox模型生存概率低于50%的对应月份

假设你的生存概率DataFrame命名为survival_probs，行索引代表月份，每一列对应一位客户的月度生存概率。以下是快速实现需求的代码方案：

核心代码 #

import pandas as pd

# 生成布尔矩阵，标记生存概率低于0.5的位置
below_threshold = survival_probs < 0.5

# 提取每个客户首次触发阈值的月份索引
first_below_half = below_threshold.idxmax(axis=0)

# 处理始终未低于0.5的客户，标记为缺失值或自定义文本
first_below_half[~below_threshold.any(axis=0)] = pd.NA

# 转换为易读的结果格式
result_df = first_below_half.reset_index()
result_df.columns = ["客户ID", "首次生存概率<0.5的月份"]

代码解释 #

• below_threshold = survival_probs < 0.5：创建与原DataFrame结构一致的布尔矩阵，True表示对应月份该客户的生存概率跌破0.5。
• below_threshold.idxmax(axis=0)：对每一列（客户）取第一个True值的行索引——因为布尔值中True等价于1，False等价于0，idxmax会返回第一个最大值的位置，正好匹配"首次低于0.5"的需求。
• 第三行代码专门处理那些所有月份生存概率都≥0.5的客户，你可以把pd.NA替换成类似"未达到阈值"的自定义文本。

示例输出 #

假设原数据如下：

运行代码后得到的结果：

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