Sid
Sid

Reputation: 286

Creating this numpy array in Python

I have the following numpy array

import numpy as np
a = np.array([1,2,6,8])

I want to create another numpy array from a such that it contains all the different possible sums of TWO elements of a. It's easy to show then that there are int(a.size*(a.size-1)/2) different possible sums, composed from:

a[0] + a[1]
a[0] + a[2]
a[0] + a[3]
a[1] + a[2]
a[1] + a[3]
a[2] + a[3]

How can I construct a numpy array with the above sums as elements without using a double for loop (the only way I can think of it). For the above example, the output should be [3,7,9,8,10,14]

MWE

eff = int(a.size*(a.size-1)/2)
c = np.empty((0, eff))

Upvotes: 2

Views: 117

Answers (3)

Paul Panzer
Paul Panzer

Reputation: 53029

You can use triu_indices:

i0,i1 = np.triu_indices(4,1)
a[i0]
# array([1, 1, 1, 2, 2, 6])
a[i1]
# array([2, 6, 8, 6, 8, 8])
a[i0]+a[i1]
# array([ 3,  7,  9,  8, 10, 14])

For more terms we need to build our own "nd_triu_idx". Here is how to do it for 3 terms out of a list of 5:

n = 5
full = np.mgrid[:n,:n,:n]
nd_triu_idx = full[:,(np.diff(full,axis=0)>0).all(axis=0)]
nd_triu_idx
# array([[0, 0, 0, 0, 0, 0, 1, 1, 1, 2],
#        [1, 1, 1, 2, 2, 3, 2, 2, 3, 3],
#        [2, 3, 4, 3, 4, 4, 3, 4, 4, 4]])

To fully generalize the number of terms use something like

k = 4
full = np.mgrid[k*(slice(n),)]

etc.

Upvotes: 5

keepAlive
keepAlive

Reputation: 6655

What about computing the cartesian product of exponentiated version of a?

>>> a = np.array([1, 2, 6, 8])[:, None]
>>> b = np.exp(a)
>>> np.unique(np.tril(np.log(np.dot(b, b.T)), k=-1))[1:]
array([ 3.,  7.,  8.,  9., 10., 14.])

Upvotes: 0

Austin
Austin

Reputation: 26039

You can do combinations on your array of size 2 and sum each one:

import numpy as np
from itertools import combinations

a = np.array([1,2,6,8])

print(list(map(sum, combinations(a, 2))))
# [3, 7, 9, 8, 10, 14]

Or using numpy:

import numpy as np

a = np.array([1,2,6,8,1])

b = a + a[:,None]
print(b[np.triu_indices(4, 1)])
# [ 3  7  9  8 10 14]

Upvotes: 2

Related Questions