Sama
Sama

Reputation: 7

Splitting a nump array at specific locations

Hey so i basically have a problem like this:

i have a numpy array which contains a matrix of values, for example:

Data = np.array([
    [3, 0, 1, 5],
    [0, 0, 0, 7],
    [0, 3, 0, 0],
    [0, 0, 0, 6],
    [5, 1, 0, 0]])

Using another array i want to extract the specific values and sum them together, this is a bit hard to explain so ill just show an example:

values = np.array([3,1,3,4,2])

so this means we want the first 3 values of the first row, first value of the second row, first 3 values of the 3rd row, first 4 values of the 4th row and first 2 values of the the last row, so we only want this data:

final_data = np.array([
    [3, 0, 1],
    [0],
    [0, 3, 0],
    [0, 0, 0, 6],
    [5, 1]])

then we want to get the sum amount of those values, in this case the sum value will be 19.

Is there any easy way to do this? also, the data isn't always the same size so i cant have any fixed variables.

Upvotes: 0

Views: 58

Answers (3)

Bobby Ocean
Bobby Ocean

Reputation: 3328

An even better answer:

Data[np.arange(Data.shape[1])<values[:,None]].sum()

Upvotes: 1

Bobby Ocean
Bobby Ocean

Reputation: 3328

You can accomplish this with advanced indexing. The advanced coordinates can be calculated separately before pulling them from the array.

Explicitly:

Data = np.array([
    [3, 0, 1, 5],
    [0, 0, 0, 7],
    [0, 3, 0, 0],
    [0, 0, 0, 6],
    [5, 1, 0, 0]])
values = np.array([3,1,3,4,2])
X = [0,0,0,1,2,2,2,3,3,3,3,4,4]
Y = [0,1,2,0,0,1,2,0,1,2,3,0,1]
Data[X,Y]

Notice X is the number of times to access each row and Y is the column to access with each X. These can be calculated from values directly:

X = np.concatenate([[n]*i for n,i in enumerate(values)])
Y = np.concatenate([np.arange(i) for i in values])

Upvotes: 0

bb1
bb1

Reputation: 7863

You can try:

sum([Data[i, :j].sum() for i, j in enumerate(values)])

Upvotes: 0

Related Questions