Search code examples
pythonmultidimensional-arrayenumeration

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?


Solution

  • 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)