Panchi
Panchi

Reputation: 551

Python : summation over all permutations

I am stuck with a seemingly easy problem, could someone please help?

I have two lists a and b. I can call the elements of the list as a[i][j] where 0<i<100 and 0<j<100.

I want to find all the a[i][j] - b[k][l], where 0<i<100, 0<j<100, 0<k<100 and 0<l<100. And then do summation over all permutations of i, j, k, l.

Anyone knows an elegant simple way to this??

Upvotes: 5

Views: 1311

Answers (3)

Andr&#233; Laszlo
Andr&#233; Laszlo

Reputation: 15537

This is way faster:

(sum(map(sum, a)) - sum(map(sum, b))) * len(a) * len(b)

Since you are calculating the sum of the total difference of the arrays. It will work in O(n), but Kevin's answer using itertools is in O(n^2).

I haven't proven it (edit: here's a proof outline), only tested using two random 100x100 arrays, but it's kind of intuitive if you think about it.

The code above will, like Kevin pointed out, assume that your arrays are 100x100 and will include 0. I assumed that you simply wanted to use two 100x100 arrays and got the notation a bit messed up - if you actually want to skip 0 and don't know the size of the array, then you can use the following snippet instead:

from itertools import product

C = 100

def _sum(arr, a, b):
    return sum(arr[i][j] for i, j in itertools.product(range(a, b), repeat=2))

answer = (_sum(a, 1, C) - _sum(b, 1, C)) * (C-1)**2

Not as pretty, since we are iterating over a submatrix, but still fast. If you want to use numpy, the sliced version can be written a little bit more concisely:

import numpy as np
A = np.array(a, np.int32)
B = np.array(b, np.int32)
answer = (np.sum(A[1:100, 1:100]) - np.sum(B[1:100, 1:100])) * 99**2

Upvotes: 2

FuzzyDuck
FuzzyDuck

Reputation: 1521

Try generating the Cartesian produce of the array ranges to get the indices.

import itertools

cprod = itertools.product(range(1, 100, 1), range(1, 100, 1), range(1, 100, 1), range(1, 100, 1))
result = sum([a[cp[0]][cp[1] - b[cp[2]][cp[3] for cp in cprod])

Upvotes: 0

Kevin
Kevin

Reputation: 30191

import itertools
sum(a[i][j] - b[k][l] for i, j, k, l in itertools.product(range(1, 100), repeat=4))

itertools.product is equivalent to a nested for loop. It will run over every (i, j, k, l) tuple from (1, 1, 1, 1) to (99, 99, 99, 99). This will skip zero, which you seem to have asked for.

Upvotes: 3

Related Questions