How index results by ID in array output python?-CodePudding

I have the following calculation below, I am trying to get the results , with their cust_id for each customer.

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

desired output example: x111([[ 55, 56, 57, 58, 59],...

CodePudding user response：

This is what I was looking for:

Dict_res = dict(zip(cust_id, results))

CodePudding user response：

If that’s what you want, you can create a dataframe with columns that are the rows of the results, and set the column names as cust_id.

pd.DataFrame(results.T, columns=cust_id.tolist())