find-demo.py

find-demo.py

"""
keys:
[3, 4]

r (10, 5):
[[22.  7. 20. 21.  6.]
 [ 8. 14.  6.  5.  5.]
 [ 1.  0.  1.  5.  4.]
 [17. 15.  0.  7.  6.]
 [ 0.  7.  7.  7.  4.]
 [23.  1.  0.  0.  1.]
 [ 2.  4. 25.  5.  5.]
 [16.  2.  5. 10.  9.]
 [ 7. 16. 11. 12.  5.]
 [ 3.  4. 17.  9.  6.]]

mask:
[False False  True False  True False  True False False  True]

rows (4, 5):
[[ 1.  0.  1.  5.  4.]
 [ 0.  7.  7.  7.  4.]
 [ 2.  4. 25.  5.  5.]
 [ 3.  4. 17.  9.  6.]]
"""


import numpy as np

if __name__ == '__main__':
  # Config
  n, size = 10, 5
  keys = [3, 4]
  
  # Create fake data
  r = np.random.randn(n * size).reshape(n, size)
  r = np.abs((r*10).round())

  # Create a mask for rows that contain the key.
  mask = np.isin(r, keys).any(1)

  # Get the relevent rows.
  rows = r[mask]

  print('keys:\n{}\n'.format(keys))
  print('r {}:\n{}\n'.format(r.shape, r))
  print('mask:\n{}\n'.format(mask))
  print('rows {}:\n{}\n'.format(rows.shape, rows))