numpy寻找数组中符合的index

import numpy as np

if __name__ == '__main__':

    a = np.array([[1, 2],
                  [3, 4]])
    pos = np.where(a > 3)
    print(pos)
    b = a[pos]
    print(b)

    pos = np.where(a<3)
    print(pos)
    b = a[pos]
    print(b)

ans:

(array([1], dtype=int64), array([1], dtype=int64))
[4]
(array([0, 0], dtype=int64), array([0, 1], dtype=int64))
[1 2]