Pythonic algorithm for doing an array calculation

Question

I am looking for the most efficient and pythonic algorithm for doing an array calculation. Here is the problem:

I have an array of shape (5,2,3) and its sum along the axis=0 as follows:

import numpy as np
A = np.array([[[ 6, 15, 89],
                 [49, 62, 12]],

                [[92,  8, 34],
                 [93, 81, 35]],

                [[ 8, 35, 63],
                 [68, 89,  5]],

                [[27, 20, 85],
                 [87, 42, 90]],

                [[99, 64, 12],
                 [90, 93, 87]]])

B = A.sum(axis=0)

So B is basically equal to A[0]+A[1]+A[2]+A[3]+A[4] which is:

array([[232, 142, 283],
       [387, 367, 229]])

I want to know at what stage of the sum process, each of 6 elements of B has gone bigger than 100. For example element B[0,0] goes above 100 after 3 steps: A[0]+A[1]+A[2], or B[1,1] goes above 100 after 2 steps A[0]+A[1]. So the final output of the algorithm should be this array:

array([[3, 5, 2],
       [2, 2, 4]])

I know I can do the calculation for each element separately but I was wondering if anyone could come up with a creative and faster algorithm.

Cheers,

Divakar · Accepted Answer · 2017-03-05 18:31:09Z

4

Use cumsum to get a cumulative summation, compare it against the threshold and finally use argmax to catch it as the first instance of crossing that threshold -

(A.cumsum(axis=0) > 100).argmax(axis=0)+1

Sample run -

In [228]: A
Out[228]: 
array([[[ 6, 15, 89],
        [49, 62, 12]],

       [[92,  8, 34],
        [93, 81, 35]],

       [[ 8, 35, 63],
        [68, 89,  5]],

       [[27, 20, 85],
        [87, 42, 90]],

       [[99, 64, 12],
        [90, 93, 87]]])

In [229]: (A.cumsum(0) > 100).argmax(0)+1
Out[229]: 
array([[3, 5, 2],
       [2, 2, 4]])

answered Mar 5, 2017 at 18:31

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Pythonic algorithm for doing an array calculation

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related