Pythonic way of iterating over 3D array

前端 未结 2 1025
庸人自扰
庸人自扰 2020-12-14 13:12

I have a 3D array in Python and I need to iterate over all the cubes in the array. That is, for all (x,y,z) in the array\'s dimensions I need to access the cub

相关标签:
2条回答
  • 2020-12-14 13:39

    Have a look at itertools, especially itertools.product. You can compress the three loops into one with

    import itertools
    
    for x, y, z in itertools.product(*map(xrange, (x_dim, y_dim, z_dim)):
        ...
    

    You can also create the cube this way:

    cube = numpy.array(list(itertools.product((0,1), (0,1), (0,1))))
    print cube
    array([[0, 0, 0],
           [0, 0, 1],
           [0, 1, 0],
           [0, 1, 1],
           [1, 0, 0],
           [1, 0, 1],
           [1, 1, 0],
           [1, 1, 1]])
    

    and add the offsets by a simple addition

    print cube + (10,100,1000)
    array([[  10,  100, 1000],
           [  10,  100, 1001],
           [  10,  101, 1000],
           [  10,  101, 1001],
           [  11,  100, 1000],
           [  11,  100, 1001],
           [  11,  101, 1000],
           [  11,  101, 1001]])
    

    which would to translate to cube + (x,y,z) in your case. The very compact version of your code would be

    import itertools, numpy
    
    cube = numpy.array(list(itertools.product((0,1), (0,1), (0,1))))
    
    x_dim = y_dim = z_dim = 10
    
    for offset in itertools.product(*map(xrange, (x_dim, y_dim, z_dim))):
        work_with_cube(cube+offset)
    

    Edit: itertools.product makes the product over the different arguments, i.e. itertools.product(a,b,c), so I have to pass map(xrange, ...) with as *map(...)

    0 讨论(0)
  • 2020-12-14 13:39
    import itertools
    for x, y, z in itertools.product(xrange(x_size), 
                                     xrange(y_size), 
                                     xrange(z_size)):
        work_with_cube(array[x, y, z])
    
    0 讨论(0)
提交回复
热议问题