如何在Python中按ID对数组输出的计算结果进行索引关联？

你的代码已经完成了核心计算逻辑，但缺少将计算结果与cust_id关联输出的步骤。以下是直接解决需求的方案：

原代码回顾 #

import numpy as np
import pandas as pd

matrix_data = np.arange(125).reshape(5,5,5)

customer_data = pd.DataFrame({"cust_id": ['x111', 'x222', 'x333'],
                              "state": [2, 3, 2],
                              "amount": [1, 2, 5],
                              "year": [2, 3, 4]})

state = customer_data['state'] -1
amount = customer_data['amount'].to_numpy().reshape(-1, 1)
year = customer_data['year']
cust_id=customer_data['cust_id']
results = matrix_data[year, state] * amount

修改后代码（关联cust_id输出） #

通过zip将客户ID与计算结果绑定，即可直接按ID输出对应内容：

# 生成客户ID与结果的配对字典
cust_result_map = dict(zip(cust_id, results))

# 逐个输出客户ID及对应计算结果
for id, res in cust_result_map.items():
    print(f"{id}: {res}")

中文期望输出示例 #

运行上述代码后，输出格式如下：

x111: [[55 56 57 58 59]
[60 61 62 63 64]
[65 66 67 68 69]
[70 71 72 73 74]
[75 76 77 78 79]]
x222: [[160 162 164 166 168]
[170 172 174 176 178]
[180 182 184 186 188]
[190 192 194 196 198]
[200 202 204 206 208]]
x333: [[525 530 535 540 545]
[550 555 560 565 570]
[575 580 585 590 595]
[600 605 610 615 620]
[625 630 635 640 645]]

补充说明 #

• 原代码中matrix_data[year, state]会根据每个客户的year和减1后的state提取对应5×5矩阵，再通过amount完成数值缩放
• 用zip(cust_id, results)可以快速建立客户ID与结果的一一对应关系，方便后续输出或存储

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