searchsorted.nim

import algorithm

func cumsum[T](x: openarray[T]): seq[T] =
  result = newSeq[T](x.len)
  result[0] = x[0]
  for i in 1 ..< x.len:
    result[i] = x[i] + result[i-1]

func searchsorted[T](a, v: openarray[T], leftSide: static bool = true): seq[T] =
  result = newSeqUninitialized[T](v.len)
  for i in 0 ..< v.len:
    when leftSide:
      result[i] = a.lowerBound(v[i])
    else:
      result[i] = a.upperBound(v[i])

when isMainModule:
  block: # Numpy
    let p = [1, 2, 3, 4, 5]

    doAssert searchsorted(p, [3]) == @[2]
    doAssert searchsorted(p, [3], leftSide = false) == @[3]
    doassert searchsorted(p, [-10, 10, 2, 3]) == @[0, 5, 1, 2]

  block: # Tensorflow
    let row1 = [0, 3, 9, 9, 10]
    let row2 = [1, 2, 3, 4, 5]

    let val1 = [2, 4, 9]
    let val2 = [0, 2, 6]

    doAssert searchsorted(row1, val1) == @[1, 2, 2]
    doAssert searchsorted(row2, val2) == @[0, 1, 5]

    doAssert searchsorted(row1, val1, leftSide = false) == @[1, 2, 4]
    doAssert searchsorted(row2, val2, leftSide = false) == @[0, 2, 5]