Kth Smallest Element [Medium]

https://www.interviewbit.com/problems/kth-smallest-element-in-the-array/

Straightforward problem description: Find the kth smallest element in an array. Easiest approach I could think of is to use a queue. Didn't see a more optimal solution on the platform (InterviewBit) that this question popped up on

import scala.collection.mutable.PriorityQueue

def kthsmallest(A: Array[Int], B: Int): Int = {

  @tailrec
  def go(pq: PriorityQueue[Int], i: Int): Int = {
    if (i == A.length) {
      pq.head
    } else {
      go(
        if (pq.length == B) {
          if (A(i) < pq.head) (pq += A(i)).tail
          else pq
        } else {
          pq += A(i)
        },
        i + 1
      )
    }
  }

  val order = new Ordering[Int] {
    override def compare(x: Int, y: Int): Int =
      implicitly[Ordering[Int]].compare(x, y)
  }

  A.lift(0).fold(0)(elem => go(PriorityQueue[Int](elem)(order), 1))

}