Iterate over the output of `np.where`

后端 未结 2 1097
误落风尘
误落风尘 2021-01-12 08:52

I have a 3D array and use np.where to find elements that meet a certain condition. The output of np.where is a tuple of three 1D arrays, each givin

相关标签:
2条回答
  • 2021-01-12 09:14

    Use np.transpose over zip, it's faster for large arrays

    import numpy as np
    myarray = np.random.randint(0, 7, size=1000000)
    %timeit indices = zip(*np.where(myarray == 0))
    %timeit indices = np.transpose(np.where(myarray == 0))
    
    10 loops, best of 3: 31.8 ms per loop
    100 loops, best of 3: 15.9 ms per loop
    
    0 讨论(0)
  • 2021-01-12 09:24

    Use zip

    indices = zip(*np.where(myarray == 0))
    

    Then you can do

    for i, j, k in indices:
        print ...
    

    For example,

    In [1]: x = np.random_integers(0, 1, (3, 3, 3))
    In [2]: np.where(x) # you want np.where(x==0)
    Out[2]: (array([0, 0, 0, 0, 0, 1, 1, 1, 1, 2]),
             array([0, 1, 1, 2, 2, 0, 0, 1, 1, 2]),
             array([1, 0, 1, 0, 1, 1, 2, 0, 2, 2]))
    In [3]: zip(*np.where(x))
    Out[3]: [(0, 0, 1),
             (0, 1, 0),
             (0, 1, 1),
             (0, 2, 0),
             (0, 2, 1),
             (1, 0, 1),
             (1, 0, 2),
             (1, 1, 0),
             (1, 1, 2),
             (2, 2, 2)]
    
    0 讨论(0)
提交回复
热议问题