Fix Qtree quantileBounds off-by-one error

algebird-core/src/main/scala/com/twitter/algebird/QTree.scala

@@ -43,20 +43,20 @@ object QTree {
    * level gives a bin size of 2^level. By default the bin size is 1/65536 (level = -16)
    */
   def apply[A](kv: (Double, A), level: Int = DefaultLevel): QTree[A] =
-    QTree(math.floor(kv._1 / math.pow(2.0, level)).toLong,
-      level,
-      1,
-      kv._2,
-      None,
-      None)
+    QTree(offset = math.floor(kv._1 / math.pow(2.0, level)).toLong,
+      level = level,
+      count = 1,
+      sum = kv._2,
+      lowerChild = None,
+      upperChild = None)
 
   def apply[A](kv: (Long, A)): QTree[A] =
-    QTree(kv._1,
-      0,
-      1,
-      kv._2,
-      None,
-      None)
+    QTree(offset = kv._1,
+      level = 0,
+      count = 1,
+      sum = kv._2,
+      lowerChild = None,
+      upperChild = None)
 
   /**
    * The common case of wanting an offset and sum for the same value
@@ -197,36 +197,38 @@ case class QTree[A](
   def quantileBounds(p: Double): (Double, Double) = {
     require(p >= 0.0 && p <= 1.0, "The given percentile must be of the form 0 <= p <= 1.0")
 
-    val rank = math.floor(count * p).toLong
-    // get is safe below, because findRankLowerBound only returns
-    // None if rank > count, but due to construction rank <= count
-    (findRankLowerBound(rank).get, findRankUpperBound(rank).get)
+    // rank is 0-indexed and 0 <= rank < count
+    val rank = math.floor(count * p).toLong.min(count - 1)
+
+    findRankBounds(rank)
   }
 
-  private def findRankLowerBound(rank: Long): Option[Double] =
-    if (rank > count)
-      None
-    else {
-      val childCounts = mapChildrenWithDefault(0L)(_.count)
-      val parentCount = count - childCounts._1 - childCounts._2
-      lowerChild.flatMap { _.findRankLowerBound(rank - parentCount) }
-        .orElse {
-          val newRank = rank - childCounts._1 - parentCount
-          if (newRank <= 0)
-            Some(lowerBound)
-          else
-            upperChild.flatMap{ _.findRankLowerBound(newRank) }
-        }
-    }
+  /**
+   * Precondition: if 0 <= rank < count
+   */
+  private def findRankBounds(rank: Long): (Double, Double) = {
+    require(0 <= rank && rank < count)
 
-  private def findRankUpperBound(rank: Long): Option[Double] = {
-    if (rank > count)
-      None
-    else {
-      lowerChild.flatMap{ _.findRankUpperBound(rank) }.orElse {
-        val lowerCount = lowerChild.map{ _.count }.getOrElse(0L)
-        upperChild.flatMap{ _.findRankUpperBound(rank - lowerCount) }.orElse(Some(upperBound))
-      }
+    val (leftCount, rightCount) = mapChildrenWithDefault(0L)(_.count)
+    val parentCount = count - leftCount - rightCount
+
+    if (0 <= rank && rank < leftCount) {
+
+      // lowerChild.get is safe because
+      // leftCount > 0, so lowerChild is not None.
+      lowerChild.get.findRankBounds(rank)
+
+    } else if (leftCount <= rank && rank < leftCount + parentCount) {
+      (lowerBound, upperBound)
+    } else {
+      // so leftCount + parentCount <= rank < count
+
+      // upperChild.get is safe because
+      // leftCount + parentCount < count,
+      // so leftCount + (count - leftCount - rightCount) < count,
+      // so count - rightCount < count,
+      // so rightCount > 0, so upperChild is not None.
+      upperChild.get.findRankBounds(rank - leftCount - parentCount)
     }
   }
 

algebird-test/src/test/scala/com/twitter/algebird/QTreeTest.scala

@@ -68,6 +68,21 @@ class QTreeTest extends WordSpec with Matchers {
   def trueRangeSum(list: Seq[Double], from: Double, to: Double) =
     list.filter{ _ >= from }.filter{ _ < to }.sum
 
+  for (k <- Seq(3, 11, 51, 101)) {
+    s"QTree with elements (1 to $k)" should {
+      val trueMedian = (1 + k) / 2
+      s"have median $trueMedian" in {
+        implicit val sg = new QTreeSemigroup[Double](k)
+
+        val list = (1 to k).map(_.toDouble)
+        val qtree = sg.sumOption(list.map(QTree(_))).get
+
+        val (lower, upper) = qtree.quantileBounds(0.5)
+        assert(lower <= trueMedian && trueMedian <= upper)
+      }
+    }
+  }
+
   for (k <- (1 to 6))
     ("QTree with sizeHint 2^" + k) should {
       "always contain the true quantile within its bounds" in {