How to split results of a multi-dimensional array python

Question

I have a for loop that multiplies customer data by an array list based on age and state. The loop works but I need to split the results of the output by customer for my next calculation, any Idea how to do this each customers output varies in length based on age.Example customer us111 has an output of size (12,10,10), and customer us112 has an output of size (11,10,10)

split=np.random.rand(15, 10, 10)
age = [3,4,9]
customers = np.array([
    [0,1000,0,0,0,0,0,0,0,0],
    [0,0,0,4890,0,0,0,0,0,0],
    [0,0,0,0,624,0,0,0,0,0],
])
results = []
for c, y in zip(customers, age):
    for m in split[y:]:
        c = c @ m
        results.append(c)
results = np.stack(results)

For this example split is the array list with size (15,10,10), where age is represented by each matrix example split[0]= age 1, split[1] =2 and so on until age 15 is reached.Each matrix has 10 states (AJ) for this case.

customer data below:

user_id | current_state | age   | amount
us111   |   B           |   3   |  1000
us112   |   D           |   4   |  4890
us113   |   E           |   9   |  624

The loop puts each users data in order of their current state example us111= np.array([0,1000,0,0,0,0,0,0,0,0])

and then multiplies by the matrices from the customers age until age 15 is reached.

The current ouptut of the loop appends each customers results below each other but I need to add their user_id next to each array, it can be in a dataframe as well, it doesnt have to be in an array.

Answer 1

Maybe you can separate it in a dictionary changing the loop:

   results_dict = {}
    for _id, c, y in zip(user_id ,customers, age):
        results = []
        for m in split[y:]:
            c = c @ m
            results.append(c)
    results_dict[_id] = results