list comprehension variable assignment

Question

I am attempting to create a 4d array and assign variables to each of the cells.

Typically I would use four "for loops" but this is very messy and takes up a lot of space.

What i'm currently doing:

for x in range(2):
    for y in range(2):
        for j in range(2):
            for k in range(2):
                array[x,y,j,k] = 1 #will be a function in reality

I've tried using list comprehension but this only creates the list and does not assign variables to each cell.

Are there space-efficient ways to run through multiple for loops and assign variables with only a few lines of code?

"I've tried using list comprehension but this only creates the list and does not assign variables to each cell." - a list comprehension can do everything your quadruply-nested loop is doing. Either you screwed up your comprehension, or you don't quite understand what your loops are doing. — user2357112, Mar 16 '19 at 18:11
When you say "array", what exactly do you mean? A numpy array? A list of lists? — Aran-Fey, Mar 16 '19 at 18:15
List comprehensions are for *lists*. `array` would appear to be a NumPy array (or possibly a `dict`). — chepner, Mar 16 '19 at 18:15
From what i've seen list comprehension is mostly used for initialising a list, but is NOT used for assigning variables to list indices- am i wrong about his? — Christian Potts, Mar 16 '19 at 18:25
I see, so it appears I've gotten the terms mixed up.... that would explain why it was so damn tough to find a satisfactory answer. — Christian Potts, Mar 16 '19 at 18:31
If you actually have a numpy array, then you'd not want to use a list comprehension *or* `itertools`, as [numpy arrays support iteration, with and without indices, and assignment to boot](https://docs.scipy.org/doc/numpy/reference/arrays.nditer.html). — Martijn Pieters, Mar 16 '19 at 22:15

Aran-Fey · Accepted Answer · 2019-03-16T18:17:52.400

2

Assuming you've already created an empty (numpy?) array, you can use itertools.product to fill it with values:

import itertools

for x, y, j, k in itertools.product(range(2), repeat=4):
    arr[x,y,j,k] = 1

If not all of the array's dimensions are equal, you can list them individually:

for x, y, j, k in itertools.product(range(2), range(2), range(2), range(2)):
    arr[x,y,j,k] = 1

edited Mar 16 '19 at 18:17

answered Mar 16 '19 at 18:12

Aran-Fey

39,665
11
104
149

This is exactly what I am looking for! I've searched far to long to find this answer! out of curiosity, is this actually called list comprehension or something else? – Christian Potts Mar 16 '19 at 18:19
Wouldn't [`numpy.indices()`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.indices.html) be more appropriate? – Martijn Pieters Mar 16 '19 at 22:10
Also, numpy arrays can be iterated over directly [complete with indices and a writeable reference](https://docs.scipy.org/doc/numpy/reference/arrays.nditer.html). – Martijn Pieters Mar 16 '19 at 22:12
@MartijnPieters Well, I wasn't sure if it's really a numpy array, so... – Aran-Fey Mar 16 '19 at 23:24

Mulan · Answer 2 · 2019-03-16T19:25:12.093

You may however be wondering how itertools.product does the trick. Or maybe you want to encode a different transformation in your recursive expansion. Below, I'll share one possible solution using Python's generators -

def product (*iters):
  def loop (prod, first = [], *rest):
    if not rest:
      for x in first:
        yield prod + (x,)
    else:
      for x in first:
        yield from loop (prod + (x,), *rest)
  yield from loop ((), *iters)

for prod in product ("ab", "xyz"):
  print (prod)

# ('a', 'x')
# ('a', 'y')
# ('a', 'z')
# ('b', 'x')
# ('b', 'y')
# ('b', 'z')

Because product accepts a a list of iterables, any iterable input can be used in the product. They can even be mixed as demonstrated here -

print (list (product (['@', '%'], range (2), "xy")))
# [ ('@', 0, 'x')
# , ('@', 0, 'y')
# , ('@', 1, 'x')
# , ('@', 1, 'y')
# , ('%', 0, 'x')
# , ('%', 0, 'y')
# , ('%', 1, 'x')
# , ('%', 1, 'y')
# ]

We could make a program foo that provides the output posted in your question -

def foo (n, m):
  ranges = [ range (m) ] * n
  yield from product (*ranges)

for prod in foo (4, 2):
  print (prod)

# (0, 0, 0, 0)
# (0, 0, 0, 1)
# (0, 0, 1, 0)
# (0, 0, 1, 1)
# (0, 1, 0, 0)
# (0, 1, 0, 1)
# (0, 1, 1, 0)
# (0, 1, 1, 1)
# (1, 0, 0, 0)
# (1, 0, 0, 1)
# (1, 0, 1, 0)
# (1, 0, 1, 1)
# (1, 1, 0, 0)
# (1, 1, 0, 1)
# (1, 1, 1, 0)
# (1, 1, 1, 1)

Or use destructuring assignment to create bindings for individual elements of the product. In your program, simply replace print with your real function -

for (w, x, y, z) in foo (4, 2):
  print ("w", w, "x", x, "y", y, "z", z)

# w 0 x 0 y 0 z 0
# w 0 x 0 y 0 z 1
# w 0 x 0 y 1 z 0
# w 0 x 0 y 1 z 1
# w 0 x 1 y 0 z 0
# w 0 x 1 y 0 z 1
# w 0 x 1 y 1 z 0
# w 0 x 1 y 1 z 1
# w 1 x 0 y 0 z 0
# w 1 x 0 y 0 z 1
# w 1 x 0 y 1 z 0
# w 1 x 0 y 1 z 1
# w 1 x 1 y 0 z 0
# w 1 x 1 y 0 z 1
# w 1 x 1 y 1 z 0
# w 1 x 1 y 1 z 1

Because product is defined as a generator, we are afforded much flexibility even when writing more complex programs. Consider this program that finds right triangles made up whole numbers, a Pythagorean triple. Also note that product allows you to repeat an iterable as input as see in product (r, r, r) below

def is_triple (a, b, c):
  return a * a + b * b == c * c

def solver (n):
  r = range (1, n)
  for p in product (r, r, r):
    if is_triple (*p):
      yield p

print (list (solver (20)))
# (3, 4, 5)
# (4, 3, 5)
# (5, 12, 13)
# (6, 8, 10)
# (8, 6, 10)
# (8, 15, 17)
# (9, 12, 15)
# (12, 5, 13)
# (12, 9, 15)
# (15, 8, 17)

For additional explanation and a way to see how to do this without using generators, view this answer.

list comprehension variable assignment

2 Answers2