I’ve always thought that persistent segment trees, also known as “functional segment trees” in Chinese, are more elegant than traditional ones, since I can essentially treat them as immutable objects. So recently, I decided to code an implementation in Scala, a functional language.

Without further ado, here’s the code:

object PersistentSegTree {
  def init(l: Int, r: Int): TreeNode = {
    val mid = (l + r) >>> 1
    if (l != r) Node(0, init(l, mid), init(mid + 1, r))
    else Leaf(0)
  }

  sealed trait TreeNode {
    val sum: Int

    def modify(pos: Int, v: Int, l: Int, r: Int): TreeNode

    def queryKthSmallest(old: TreeNode, k: Int, l: Int, r: Int): Int
  }

  case class Node(sum: Int, lc: TreeNode, rc: TreeNode) extends TreeNode {
    def modify(pos: Int, v: Int, l: Int, r: Int): TreeNode = {
      val mid = (l + r) >>> 1
      if (pos <= mid) Node(sum + v, lc.modify(pos, v, l, mid), rc)
      else Node(sum + v, lc, rc.modify(pos, v, mid + 1, r))
    }

    def queryKthSmallest(old: TreeNode, k: Int, l: Int, r: Int): Int = {
      val nOld = old.asInstanceOf[Node]
      val (lSum, mid) = (lc.sum - nOld.lc.sum, (l + r) / 2)
      if (k <= lSum) lc.queryKthSmallest(nOld.lc, k, l, mid)
      else rc.queryKthSmallest(nOld.rc, k - lSum, mid + 1, r)
    }
  }

  case class Leaf(sum: Int) extends TreeNode {
    def modify(pos: Int, v: Int, l: Int, r: Int): TreeNode = Leaf(sum + v)

    def queryKthSmallest(old: TreeNode, k: Int, l: Int, r: Int): Int = l
  }
}

object Main extends App {
  var Array(n, m) = io.StdIn.readLine().split(" ").map(_.toInt)
  val arr = io.StdIn.readLine().split(" ").map(_.toInt)
  // map the array to an id ranged 0 ~ n
  val dict = arr.distinct.sorted.zipWithIndex.toMap
  // build the tree
  val roots = (0 until n).foldLeft(List(PersistentSegTree.init(1, n))) { (lis, cur) =>
    lis.head.modify(dict(arr(cur)) + 1, 1, 1, n) :: lis
  }.reverse.toSeq
  while ({ m -= 1; m >= 0 }) {
    val Array(i, j, k) = io.StdIn.readLine().split(" ").map(_.toInt)
    println(roots(j).queryKthSmallest(roots(i - 1), k, 1, n))
  }
}

As long as you know how persistent segment trees work, the code should be pretty self-explanatory.

Once I finished the code, I submitted it to MKTHNUM on SPOJ for testing, and got a TLE. It turned out that the tree building process, after optimization, took about 0.6 ~ 0.7 seconds on my laptop (CPU: Intel i5-2520M), which, unfortunately, is a bit too long.

Recalling that the JVM might need some warmup due to JIT optimizations, I added a loop that ran the building process ~20 times before the benchmark, and sure enough, the time went down to 0.1 ~ 0.2 seconds.

Of course, it’s impossible to apply the method on online judges. As such, I can sort of understand why most OJs give Java programs a longer time limit.