Apply calculation to pandas dataframe-CodePudding

I am trying to replicate the calculation below, but using my customer_data frame instead of customers as an array list. Please help.

current calculation:

matrix_data = np.arange(125).reshape(5,5,5)
customers = np.array([
    [0,1,0,0,0],
    [0,0,2,0,0],
    [0,5,0,0,0]
])
years = [2, 3, 4]

new = (customers[:,:,None] * matrix_data[years]).sum(axis=1)

customer_data dataframe:

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

CodePudding user response：

# Given:

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]})

# Solution:

state = customer_data['state'] - 1  # indices start with 0, not with 1
amount = customer_data['amount'].to_numpy().reshape(-1, 1)
year = customer_data['year']

new = matrix_data[year, state] * amount

new:

array([[ 55,  56,  57,  58,  59],
       [170, 172, 174, 176, 178],
       [525, 530, 535, 540, 545]])