CODEWITHSUNDEEP
python
user9996043
user9996043

Reputation: 229

Compute squared difference of n equal length lists of an array

I'm trying to design a function that basically computes the squared difference for equal length lists based on the combinations of those lists by the formula [n(n-1)]/2:

So for my data:

arr = np.array([[10, 20, 30, 40],  #list_1
                [11, 12, 13, 14],  #list_2
                [22, 23, 25, 27],  #list_3
                [12, 19, 24, 78]]) #list_4

I want to be doing this:

final_arr = ((arr[0] - arr[1])**2 +
             (arr[0] - arr[2])**2 +
             (arr[0] - arr[3])**2 +
             (arr[1] - arr[2])**2 +
             (arr[1] - arr[3])**2 +
             (arr[2] - arr[3])**2)

Output:

list(final_arr) = [371, 260, 616, 9155]

As you can see it's basically all 6 ([(4*3)/2] = 6) combinations of lists i.e.

[(list_1, list_2), (list_1, list_3), (list_1, list_4), (list_2, list_3), (list_2, list_4), (list_3, list_4)]

I'm trying to build a function that can do this for an array containing n amount of equal length lists.

Upvotes: 0

Views: 57

Answers (1)

deadshot
deadshot

Reputation: 9071

You can use itertools.combinations and zip

from itertools import combinations

lists = [list_1, list_2, list_3, list_4]
for x in zip(*lists):
    print(sum((c[0] - c[1]) ** 2 for c in combinations(x, 2)))

Output:

371
260
616
9155

Using numpy.apply_along_axis

fun = lambda x: sum((c[0] - c[1]) ** 2 for c in combinations(x, 2))
res = np.apply_along_axis(fun, 0, arr)
print(list(res))

Output:

[371, 260, 616, 9155]

Upvotes: 4

Related Questions