将计算应用于熊猫数据帧

Question

我正在尝试复制下面的计算，但使用我的customer_data框架而不是客户作为数组列表。请帮忙。

当前计算：

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数据帧：

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

Answer 1

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