Numpy element-wise in operation

前端 未结 1 848
半阙折子戏
半阙折子戏 2020-12-11 09:39

Suppose I have a column vector y with length n, and I have a matrix X of size n*m. I want to check for each element i in y, whether the element is in the corresponding row i

1条回答
  •  南方客
    南方客 (楼主)
    2020-12-11 10:00

    Vectorized approach using broadcasting -

    ((X == y[:,None]).any(1)).astype(int)
    

    Sample run -

    In [41]: X        # Input 1
    Out[41]: 
    array([[1, 2, 3],
           [3, 4, 5],
           [4, 3, 2],
           [2, 2, 2]])
    
    In [42]: y        # Input 2
    Out[42]: array([1, 2, 3, 4])
    
    In [43]: X == y[:,None] # Broadcasted  comparison
    Out[43]: 
    array([[ True, False, False],
           [False, False, False],
           [False,  True, False],
           [False, False, False]], dtype=bool)
    
    In [44]: (X == y[:,None]).any(1) # Check for any match along each row
    Out[44]: array([ True, False,  True, False], dtype=bool)
    
    In [45]: ((X == y[:,None]).any(1)).astype(int) # Convert to 1s and 0s
    Out[45]: array([1, 0, 1, 0])
    

    0 讨论(0)
提交回复
热议问题