Nearest keys in a SortedMap

Given a key k in a SortedMap, how can I efficiently find the largest key m that is less than or equal to k, and also the smallest key n that is greater than or equal to k. Tha开发者_开发技巧nk you.

---

Looking at the source code for 2.9.0, the following code seems about to be the best you can do

def getLessOrEqual[A,B](sm: SortedMap[A,B], bound: A): B = {
  val key = sm.to(x).lastKey
  sm(key)
}

I don't know exactly how the splitting of the RedBlack tree works, but I guess it's something like a O(log n) traversal of the tree/construction of new elements and then a balancing, presumable also O(log n). Then you need to go down the new tree again to get the last key. Unfortunately you can't retrieve the value in the same go. So you have to go down again to fetch the value.

In addition the lastKey might throw an exception and there is no similar method that returns an Option.

I'm waiting for corrections.

Edit and personal comment

The SortedMap area of the std lib seems to be a bit neglected. I'm also missing a mutable SortedMap. And looking through the sources, I noticed that there are some important methods missing (like the one the OP asks for or the ones pointed out in my answer) and also some have bad implementation, like 'last' which is defined by TraversableLike and goes through the complete tree from first to last to obtain the last element.

Edit 2

Now the question is reformulated my answer is not valid anymore (well it wasn't before anyway). I think you have to do the thing I'm describing twice for lessOrEqual and greaterOrEqual. Well you can take a shortcut if you find the equal element.

---

Scala's SortedSet trait has no method that will give you the closest element to some other element.

It is presently implemented with TreeSet, which is based on RedBlack. The RedBlack tree is not visible through methods on TreeSet, but the protected method tree is protected. Unfortunately, it is basically useless. You'd have to override methods returning TreeSet to return your subclass, but most of them are based on newSet, which is private.

So, in the end, you'd have to duplicate most of TreeSet. On the other hand, it isn't all that much code.

Once you have access to RedBlack, you'd have to implement something similar to RedBlack.Tree's lookup, so you'd have O(logn) performance. That's actually the same complexity of range, though it would certainly do less work.

Alternatively, you'd make a zipper for the tree, so that you could actually navigate through the set in constant time. It would be a lot more work, of course.

---

Using Scala 2.11.7, the following will give what you want:

scala> val set = SortedSet('a', 'f', 'j', 'z')
set: scala.collection.SortedSet[Char] = TreeSet(a, f, j, z)

scala> val beforeH = set.to('h').last
beforeH: Char = f

scala> val afterH = set.from('h').head
afterH: Char = j

Generally you should use lastOption and headOption as the specified elements may not exist. If you are looking to squeeze a little more efficiency out, you can try replacing from(...).head with keysIteratorFrom(...).head

---

Sadly, the Scala library only allows to make this type of query efficiently:

and also the smallest key n that is greater than or equal to k.

val n = TreeMap(...).keysIteratorFrom(k).next

You can hack this by keeping two structures, one with normal keys, and one with negated keys. Then you can use the other structure to make the second type of query.

val n = - TreeMap(...).keysIteratorFrom(-k).next

---

Looks like I should file a ticket to add 'fromIterator' and 'toIterator' methods to 'Sorted' trait.

---

Well, one option is certainly using java.util.TreeMap.

It has lowerKey and higherKey methods, which do excatly what you want.

---

I had a similar problem: I wanted to find the closest element to a given key in a SortedMap. I remember the answer to this question being, "You have to hack TreeSet," so when I had to implement it for a project, I found a way to wrap TreeSet without getting into its internals.

I didn't see jazmit's answer, which more closely answers the original poster's question with minimum fuss (two method calls). However, those method calls do more work than needed for this application (multiple tree traversals), and my solution provides lots of hooks where other users can modify it to their own needs.

Here it is:

import scala.collection.immutable.TreeSet
import scala.collection.SortedMap

// generalize the idea of an Ordering to metric sets
trait MetricOrdering[T] extends Ordering[T] {
  def distance(x: T, y: T): Double
  def compare(x: T, y: T) = {
    val d = distance(x, y)
    if (d > 0.0) 1
    else if (d < 0.0) -1
    else 0
  }
}

class MetricSortedMap[A, B]
  (elems: (A, B)*)
  (implicit val ordering: MetricOrdering[A])
  extends SortedMap[A, B] {

  // while TreeSet searches for an element, keep track of the best it finds
  // with *thread-safe* mutable state, of course
  private val best = new java.lang.ThreadLocal[(Double, A, B)]
  best.set((-1.0, null.asInstanceOf[A], null.asInstanceOf[B]))

  private val ord = new MetricOrdering[(A, B)] {
    def distance(x: (A, B), y: (A, B)) = {
      val diff = ordering.distance(x._1, y._1)
      val absdiff = Math.abs(diff)

      // the "to" position is a key-null pair; the object of interest
      // is the other one
      if (absdiff < best.get._1)
        (x, y) match {
          // in practice, TreeSet always picks this first case, but that's
          // insider knowledge
          case ((to, null), (pos, obj)) =>
            best.set((absdiff, pos, obj))

          case ((pos, obj), (to, null)) =>
            best.set((absdiff, pos, obj))

          case _ =>
        }

      diff
    }
  }

  // use a TreeSet as a backing (not TreeMap because we need to get
  // the whole pair back when we query it)
  private val treeSet = TreeSet[(A, B)](elems: _*)(ord)

  // find the closest key and return:
  // (distance to key, the key, its associated value)
  def closest(to: A): (Double, A, B) = {
    treeSet.headOption match {
      case Some((pos, obj)) =>
        best.set((ordering.distance(to, pos), pos, obj))
      case None =>
        throw new java.util.NoSuchElementException(
          "SortedMap has no elements, and hence no closest element")
    }

    treeSet((to, null.asInstanceOf[B]))  // called for side effects

    best.get
  }

  // satisfy the contract (or throw UnsupportedOperationException)
  def +[B1 >: B](kv: (A, B1)): SortedMap[A, B1] =
    new MetricSortedMap[A, B](
      elems :+ (kv._1, kv._2.asInstanceOf[B]): _*)
  def -(key: A): SortedMap[A, B] =
    new MetricSortedMap[A, B](elems.filter(_._1 != key): _*)
  def get(key: A): Option[B] = treeSet.find(_._1 == key).map(_._2)
  def iterator: Iterator[(A, B)] = treeSet.iterator
  def rangeImpl(from: Option[A], until: Option[A]): SortedMap[A, B] =
    new MetricSortedMap[A, B](treeSet.rangeImpl(
      from.map((_, null.asInstanceOf[B])),
      until.map((_, null.asInstanceOf[B]))).toSeq: _*)
}

// test it with A = Double
implicit val doubleOrdering =
  new MetricOrdering[Double] {
    def distance(x: Double, y: Double) = x - y
  }

// and B = String
val stuff = new MetricSortedMap[Double, String](
  3.3 -> "three",
  1.1 -> "one",
  5.5 -> "five",
  4.4 -> "four",
  2.2 -> "two")

println(stuff.iterator.toList)

println(stuff.closest(1.5))
println(stuff.closest(1000))
println(stuff.closest(-1000))
println(stuff.closest(3.3))
println(stuff.closest(3.4))
println(stuff.closest(3.2))

---

I've been doing:

val m = SortedMap(myMap.toSeq:_*)
val offsetMap = (m.toSeq zip m.keys.toSeq.drop(1)).map {
  case ( (k,v),newKey) => (newKey,v)
}.toMap

When I want the results of my map off-set by one key. I'm also looking for a better way, preferably without storing an extra map.