mask a 2D numpy array based on values in one column

前端 未结 3 743
萌比男神i
萌比男神i 2020-12-14 22:38

Suppose I have the following numpy array:

a = [[1, 5, 6],
     [2, 4, 1],
     [3, 1, 5]]

I want to mask all the rows which have 1

3条回答
  •  悲&欢浪女
    2020-12-14 22:48

    You can create the desired mask by

    mask = numpy.repeat(a[:,0]==1, a.shape[1])
    

    and the masked array by

    masked_a = numpy.ma.array(a, mask=numpy.repeat(a[:,0]==1, a.shape[1]))
    

提交回复
热议问题