Replace values of a numpy index array with values of a list

后端 未结 6 773
迷失自我
迷失自我 2020-12-15 04:31

Suppose you have a numpy array and a list:

>>> a = np.array([1,2,2,1]).reshape(2,2)
>>> a
array([[1, 2],
       [2, 1]])
>>> b = [         


        
6条回答
  •  时光说笑
    2020-12-15 05:05

    Instead of replacing the values one by one, it is possible to remap the entire array like this:

    import numpy as np
    a = np.array([1,2,2,1]).reshape(2,2)
    # palette must be given in sorted order
    palette = [1, 2]
    # key gives the new values you wish palette to be mapped to.
    key = np.array([0, 10])
    index = np.digitize(a.ravel(), palette, right=True)
    print(key[index].reshape(a.shape))
    

    yields

    [[ 0 10]
     [10  0]]
    

    Credit for the above idea goes to @JoshAdel. It is significantly faster than my original answer:

    import numpy as np
    import random
    palette = np.arange(8)
    key = palette**2
    a = np.array([random.choice(palette) for i in range(514*504)]).reshape(514,504)
    
    def using_unique():
        palette, index = np.unique(a, return_inverse=True)
        return key[index].reshape(a.shape)
    
    def using_digitize():
        index = np.digitize(a.ravel(), palette, right=True)
        return key[index].reshape(a.shape)
    
    if __name__ == '__main__':
        assert np.allclose(using_unique(), using_digitize())
    

    I benchmarked the two versions this way:

    In [107]: %timeit using_unique()
    10 loops, best of 3: 35.6 ms per loop
    In [112]: %timeit using_digitize()
    100 loops, best of 3: 5.14 ms per loop
    

提交回复
热议问题