Python – enumerate a multidimensional list

Let’s say I have a multidimensional array, foobar:

foobar = [[[0, 1, 2],
           [3, 4, 5, 6],
           [7, 8]],
          [[9, 10],
           [11, 12, 13, 14, 15],
           [16, 17, 18],
           [19, 20, 21, 22]],
          [[23, 24, 25],
           [26, 27]]]

Note that foobar is jagged.

The thing I need to do is replace each number in foobar with a tuple containing that number and its exact position in foobar. I also need to be able to do this when the number of dimensions and whether or not foobar is jagged is unknown.

Here is something similar, except it only works for 2 dimensions:

def enum_multidim(data):
    for i, a in enumerate(data):
        for j, b in enumerate(a):
            yield (i, j, b)

Is there a user-defined function that can do what I said above?

Answer

Recursive generator for an arbitrary number of dimensions

Code

def enum_multidim(data, t = None):
    if t is None:
        t = ()
    if not isinstance(data, list):
        yield t + (data,)
    else:
        for i, v in enumerate(data):
            yield from enum_multidim(v, t + (i,))

Test

for t in enum_multidim(foobar):
    print(t)

# Out:
(0, 0, 0, 0)
(0, 0, 1, 1)
(0, 0, 2, 2)
(0, 1, 0, 3)
(0, 1, 1, 4)
(0, 1, 2, 5)
(0, 1, 3, 6)
(0, 2, 0, 7)
(0, 2, 1, 8)
(1, 0, 0, 9)
(1, 0, 1, 10)
(1, 1, 0, 11)
(1, 1, 1, 12)
(1, 1, 2, 13)
(1, 1, 3, 14)
(1, 1, 4, 15)
(1, 2, 0, 16)
(1, 2, 1, 17)
(1, 2, 2, 18)
(1, 3, 0, 19)
(1, 3, 1, 20)
(1, 3, 2, 21)
(1, 3, 3, 22)
(2, 0, 0, 23)
(2, 0, 1, 24)
(2, 0, 2, 25)
(2, 1, 0, 26)
(2, 1, 1, 27)