Make opperations with the elements of an array without a for loop

Viewed 59

I have this problem in which I need to take an element of array make the sum between this element and the elements with higher index individually. I've done this with a for loop like:

sumtot = np.array([])
for j in range(0,len(matpos)-1):
    sum = matpos[j] + matpos[j+1:]
    sumtot = np.append(sumtot, sum)

But this takes a lot of computing time because the array matpos is a very big array, so I guess if there is a way to do this without a for loop.

A simple sample for this is:

-Input:

matpos = np.array([0, 1, 2, 3])

-Output

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

this is [0+1, 0+2, 0+3, 1+2, 1+3, 2+3]

Thank you all very much!

2 Answers

To get the array you want, you can use something like np.triu_indices with some additional manipulation:

r, c = np.triu_indices(len(matpos), 1)
totsum = matpos[r] + matpos[c]

This is probably as legible as you can get. If you want a one-liner, you can stack the indices and add up the result:

totsum = matpos[np.stack(np.triu_indices(len(matpos), 1))].sum(0)

Notice that the terms are summed as follows:

 [matpos[0], matpos[0], matpos[0], ...] + [matpos[1], matpos[2], matpos[3], ...]
 [matpos[1], matpos[1], ...] + [matpos[2], matpos[3], ...]
 [matpos[2], ...] + [matpos[3], ...]

The selected index and number of repetitions correspond exactly to the result of triu_indices (column off by one), which returns the indices of the upper triangle of a matrix.

An alternative formulation:

r, c = np.triu_indices(len(matpos) - 1)
totsum = matpos[r] + matpos[c + 1]

OR

totsum = matpos[np.stack(np.triu_indices(len(matpos) - 1), -1) + [0, 1]].sum(1)

You really should include a sample run - so we can see what your code produces without running it ourselves.

In [85]: matpos = np.arange(10)
In [86]: sumtot = np.array([])
    ...: for j in range(0,len(matpos)-1):
    ...:     sum = matpos[j] + matpos[j+1:]
    ...:     sumtot = np.append(sumtot, sum)
    ...: 
In [87]: sumtot
Out[87]: 
array([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.,  3.,  4.,  5.,  6.,
        7.,  8.,  9., 10.,  5.,  6.,  7.,  8.,  9., 10., 11.,  7.,  8.,
        9., 10., 11., 12.,  9., 10., 11., 12., 13., 11., 12., 13., 14.,
       13., 14., 15., 15., 16., 17.])

But that's not particularly instructive. In addition np.append is slow compared to list append.

So let's use lists:

In [88]: sumtot = []
    ...: for j in range(0,len(matpos)-1):
    ...:     sum = matpos[j] + matpos[j+1:]
    ...:     sumtot.append(sum)
    ...: 
    ...: 
In [89]: sumtot
Out[89]: 
[array([1, 2, 3, 4, 5, 6, 7, 8, 9]),
 array([ 3,  4,  5,  6,  7,  8,  9, 10]),
 array([ 5,  6,  7,  8,  9, 10, 11]),
 array([ 7,  8,  9, 10, 11, 12]),
 array([ 9, 10, 11, 12, 13]),
 array([11, 12, 13, 14]),
 array([13, 14, 15]),
 array([15, 16]),
 array([17])]

That gives a better idea of what you are doing - not that it looks particularly logical :)

Or to get a flat list, use extend:

In [90]: sumtot = []
    ...: for j in range(0,len(matpos)-1):
    ...:     sum = matpos[j] + matpos[j+1:]
    ...:     sumtot.extend(sum)
    ...: 
In [91]: sumtot
Out[91]: 
[1,
 2,
 3,
 ...
 16,
 17]

or for a prettier display:

In [92]: np.array(sumtot)
Out[92]: 
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9,  3,  4,  5,  6,  7,  8,  9, 10,
        5,  6,  7,  8,  9, 10, 11,  7,  8,  9, 10, 11, 12,  9, 10, 11, 12,
       13, 11, 12, 13, 14, 13, 14, 15, 15, 16, 17])

I suspect this list extend version is as good as we'll get. The ragged list in [89] indicates that a "pure" numpy solution is unlikely, or at best convoluted.

edit

With your new example:

In [93]: matpos = np.array([0, 1, 2, 3])
...
In [96]: sumtot = []
    ...: for j in range(0,len(matpos)-1):
    ...:     sum = matpos[j] + matpos[j+1:]
    ...:     sumtot.append(sum)
    ...: 
    ...: 
In [97]: sumtot
Out[97]: [array([1, 2, 3]), array([3, 4]), array([5])]

with triu:

Another answer use triu.

Let's experiement:

In [98]: matpos[:,None]+matpos
Out[98]: 
array([[0, 1, 2, 3],
       [1, 2, 3, 4],
       [2, 3, 4, 5],
       [3, 4, 5, 6]])
In [100]: np.tril(__,-1)
Out[100]: 
array([[0, 0, 0, 0],
       [1, 0, 0, 0],
       [2, 3, 0, 0],
       [3, 4, 5, 0]])

And for the larger example:

In [101]: matpos = np.arange(10)
In [102]: matpos[:,None]+matpos
Out[102]: 
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10],
       [ 2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
       [ 3,  4,  5,  6,  7,  8,  9, 10, 11, 12],
       [ 4,  5,  6,  7,  8,  9, 10, 11, 12, 13],
       [ 5,  6,  7,  8,  9, 10, 11, 12, 13, 14],
       [ 6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 7,  8,  9, 10, 11, 12, 13, 14, 15, 16],
       [ 8,  9, 10, 11, 12, 13, 14, 15, 16, 17],
       [ 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]])
In [103]: np.tril(_,-1)
Out[103]: 
array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 1,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 2,  3,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 3,  4,  5,  0,  0,  0,  0,  0,  0,  0],
       [ 4,  5,  6,  7,  0,  0,  0,  0,  0,  0],
       [ 5,  6,  7,  8,  9,  0,  0,  0,  0,  0],
       [ 6,  7,  8,  9, 10, 11,  0,  0,  0,  0],
       [ 7,  8,  9, 10, 11, 12, 13,  0,  0,  0],
       [ 8,  9, 10, 11, 12, 13, 14, 15,  0,  0],
       [ 9, 10, 11, 12, 13, 14, 15, 16, 17,  0]])

So the nonzero values match.

And extracting values with triu_indices:

In [112]: idx = np.triu_indices_from(Out[102],1)
In [113]: Out[102][idx]
Out[113]: 
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9,  3,  4,  5,  6,  7,  8,  9, 10,
        5,  6,  7,  8,  9, 10, 11,  7,  8,  9, 10, 11, 12,  9, 10, 11, 12,
       13, 11, 12, 13, 14, 13, 14, 15, 15, 16, 17])
Related