3
res = sum((i+j)%k == 0 for x, i in enumerate(a) for j in a[x+1:])

where a is an array.

I cannot understand what this code is doing..is the i in enumerate(a) for j in a[x+1:] basically a for loop inside that line to save space?

Also how can I write code like this? I am trying to learn python.

3 Answers3

5

This is a generator expression inside a call to the sum function. sum just adds things up, so let's look at the generator expression:

(i+j)%k == 0 for x, i in enumerate(a) for j in a[x+1:]

The (i+j)%k == 0 part is a boolean expression. It's either true if i+j has a remainder of 0 when dived by k or false otherwise. As a neat little trick, the True value has a numeric value of 1 in python, and False has a numeric value of 0. So this is basically just counting.

for x, i in enumerate(a) for j in a[x+1:]

This is essentially a nested for loop.

for x, i in enumerate(a):
    for j in a[x+1:]:

enumerate is a function that yields items from an iterable paired with their index, so [a, b, c] becomes [(0, a), (1, b), (2, c)]

Stick it all together, and this code calculates the number of pairs of elements in a list such that their sum is divisible byk

Patrick Haugh
  • 59,226
  • 13
  • 88
  • 96
3

you can "unwind" the code as follows:

counter=0   # emulates `sum`
for x, i in enumerate(a):
    for j in a[x+1:]:
       counter += (i+j)%k == 0

so it counts the occurrences of (i+j) being divisible by k, avoiding to count the same half and identical elements (only counts the upper matrix triangle)

List/generator comprehensions are better/faster and more concise when you're creating a list out of another iterable or iterable of iterable, like here: you created a list of booleans (0 or 1) that you sum to count the occurrences of True.

But don't abuse of them: proper use of list comprehensions - python

sometimes a plain loop to call n times a function is better (also when a side-effect is required).

When used unwisely they can become a nightmare to understand (specially the one-liners with side-effects in them)

Community
  • 1
  • 1
Jean-François Fabre
  • 137,073
  • 23
  • 153
  • 219
1

These are called comprehensions. enumerate iterates through an iterable object and returns a tuple (a, b) where a is the index in the iterable, and b is object at that index.

This is a nice piece of pythonic code. As for learning how to understand/write stuff like this, practice, and reading python documentation would be a great place to start. Remember, this is simply just syntactic sugar, you can do the same thing in multiple lines and it will work almost exactly the same.

Alex
  • 1,432
  • 14
  • 26