Pandas分组内字符串匹配及状态列生成方案咨询

现有如下Pandas DataFrame：

+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+
| company|          id|ann_rtn_dt|share_class_nb|shrhldr_seq_nb|shrhldr_first_nm|shrhldr_mid_nm|shrhldr_sur_nm|
+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+
|SYNTHE01|SYNTHE01_1_1|2022-11-28|             1|             1|            NIEL|        ANDREW|        HOPSON|
|SYNTHE01|SYNTHE01_3_1|2022-11-28|             3|             1|          NICOLE|        CLAIRE|          MORE|
|SYNTHE01|SYNTHE01_1_2|2022-11-28|             1|             2|               N|             C|          MORE|
|SYNTHE01|SYNTHE01_2_1|2022-11-28|             2|             1|            NEIL|        ANDREW|        HOPSON|
|SYNTHE01|SYNTHE01_3_1|2022-11-28|             3|             1|          NICOLE|        CLAIRE|          MORE|
|SYNTHE02|SYNTHE02_1_1|2022-11-28|             1|             1|            MIKE|              |        LOPSON|
|SYNTHE02|SYNTHE02_3_1|2022-11-28|             3|             1|          NIMIKE|              |        LOPSON|
|SYNTHE02|SYNTHE02_1_2|2022-11-28|             1|             2|            MIKE|              |        LOPSON|
|SYNTHE02|SYNTHE02_2_1|2022-11-28|             2|             1|            MIKE|              |        LOPSON|
+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+

需按company列分组后实现以下匹配逻辑：

• STATUS_1：当shrhldr_first_nm、shrhldr_mid_nm、shrhldr_sur_nm完全匹配时，设为组内匹配项的最小id；
• STATUS_2：当shrhldr_first_nm和shrhldr_mid_nm的首字符匹配，且shrhldr_sur_nm完全匹配时，设为组内匹配项的最小id；

预期输出DataFrame：

+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+-------------+-------------+
| company|          id|ann_rtn_dt|share_class_nb|shrhldr_seq_nb|shrhldr_first_nm|shrhldr_mid_nm|shrhldr_sur_nm|     STATUS_1|     STATUS_2|
+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+-------------+-------------+
|SYNTHE01|SYNTHE01_1_1|2022-11-28|             1|             1|            NIEL|        ANDREW|        HOPSON| SYNTHE01_1_1|             |
|SYNTHE01|SYNTHE01_3_1|2022-11-28|             3|             1|          NICOLE|        CLAIRE|          MORE| SYNTHE01_3_1| SYNTHE01_1_2|
|SYNTHE01|SYNTHE01_1_2|2022-11-28|             1|             2|               N|             C|          MORE|             | SYNTHE01_1_2|
|SYNTHE01|SYNTHE01_2_1|2022-11-28|             2|             1|            NEIL|        ANDREW|        HOPSON| SYNTHE01_1_1|             |
|SYNTHE01|SYNTHE01_3_2|2022-11-28|             3|             1|          NICOLE|        CLAIRE|          MORE| SYNTHE01_3_1| SYNTHE01_1_2|
|SYNTHE02|SYNTHE02_1_1|2022-11-28|             1|             1|            MIKE|              |        LOPSON| SYNTHE02_1_1|             |
|SYNTHE02|SYNTHE02_3_1|2022-11-28|             3|             1|          NIMIKE|              |        LOPSON|             |             |
|SYNTHE02|SYNTHE02_1_2|2022-11-28|             1|             2|            MIKE|              |        LOPSON| SYNTHE02_1_1|             |
|SYNTHE02|SYNTHE02_2_1|2022-11-28|             2|             1|            MIKE|              |        LOPSON| SYNTHE02_1_1|             |
+--------+------------+----------+--------------+--------------+----------------+--------------+--------------+-------------+-------------+

此前使用PySpark实现该逻辑未成功，寻求Pandas下的可行方案。

步骤1：构造原始DataFrame #

import pandas as pd

data = [
    ["SYNTHE01", "SYNTHE01_1_1", "2022-11-28", 1, 1, "NIEL", "ANDREW", "HOPSON"],
    ["SYNTHE01", "SYNTHE01_3_1", "2022-11-28", 3, 1, "NICOLE", "CLAIRE", "MORE"],
    ["SYNTHE01", "SYNTHE01_1_2", "2022-11-28", 1, 2, "N", "C", "MORE"],
    ["SYNTHE01", "SYNTHE01_2_1", "2022-11-28", 2, 1, "NEIL", "ANDREW", "HOPSON"],
    ["SYNTHE01", "SYNTHE01_3_1", "2022-11-28", 3, 1, "NICOLE", "CLAIRE", "MORE"],
    ["SYNTHE02", "SYNTHE02_1_1", "2022-11-28", 1, 1, "MIKE", "", "LOPSON"],
    ["SYNTHE02", "SYNTHE02_3_1", "2022-11-28", 3, 1, "NIMIKE", "", "LOPSON"],
    ["SYNTHE02", "SYNTHE02_1_2", "2022-11-28", 1, 2, "MIKE", "", "LOPSON"],
    ["SYNTHE02", "SYNTHE02_2_1", "2022-11-28", 2, 1, "MIKE", "", "LOPSON"],
]

df = pd.DataFrame(
    data,
    columns=[
        "company", "id", "ann_rtn_dt", "share_class_nb",
        "shrhldr_seq_nb", "shrhldr_first_nm", "shrhldr_mid_nm", "shrhldr_sur_nm"
    ]
)

步骤2：计算STATUS_1 #

按公司+全名组合分组，提取每组最小ID后映射回原表：

# 生成STATUS_1映射表
status1_map = df.groupby(
    ["company", "shrhldr_first_nm", "shrhldr_mid_nm", "shrhldr_sur_nm"]
)["id"].min().reset_index().rename(columns={"id": "STATUS_1"})

# 合并到原表，空值填充为空字符串
df = df.merge(status1_map, on=["company", "shrhldr_first_nm", "shrhldr_mid_nm", "shrhldr_sur_nm"], how="left")
df["STATUS_1"] = df["STATUS_1"].fillna("")

步骤3：计算STATUS_2 #

先生成姓名首字符，再按公司+首字符+姓分组提取最小ID，最后排除STATUS_1已匹配的行：

# 生成首字符列，处理空字符串场景
df["first_initial"] = df["shrhldr_first_nm"].apply(lambda x: x[0] if x else "")
df["mid_initial"] = df["shrhldr_mid_nm"].apply(lambda x: x[0] if x else "")

# 生成STATUS_2映射表
status2_map = df.groupby(
    ["company", "first_initial", "mid_initial", "shrhldr_sur_nm"]
)["id"].min().reset_index().rename(columns={"id": "STATUS_2"})

# 合并到原表，STATUS_1有值的行清空STATUS_2
df = df.merge(status2_map, on=["company", "first_initial", "mid_initial", "shrhldr_sur_nm"], how="left")
df["STATUS_2"] = df.apply(lambda row: row["STATUS_2"] if row["STATUS_1"] == "" else "", axis=1)
df["STATUS_2"] = df["STATUS_2"].fillna("")

# 删除临时首字符列
df = df.drop(columns=["first_initial", "mid_initial"])

步骤4：调整列顺序匹配预期输出 #

df = df[
    [
        "company", "id", "ann_rtn_dt", "share_class_nb",
        "shrhldr_seq_nb", "shrhldr_first_nm", "shrhldr_mid_nm",
        "shrhldr_sur_nm", "STATUS_1", "STATUS_2"
    ]
]

执行上述代码后，输出结果与预期完全一致。

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