CODEWITHSUNDEEP
scala
Timo Westkämper
Timo Westkämper

Reputation: 22190

Filter Map by key set

Is there a shortcut to filter a Map keeping only the entries where the key is contained in a given Set?

Here is some example code

scala> val map = Map("1"->1, "2"->2, "3"->3)
map: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2, 3 -> 3)

scala> map.filterKeys(Set("1","2").contains)
res0: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2)

I am searching for something shorter than this.

Upvotes: 29

Views: 20308

Answers (3)

missingfaktor
missingfaktor

Reputation: 92036

A tangential tip, in case you are going to follow the PredicateW idea in @oxbow_lakes' answer:

In functional programming, instead of defining ad hoc functions, we aim for more generalized and composable abstractions. For this particular case, Applicative fits the bill.

Set themselves are functions, and the Applicative instance for [B]Function1[A, B] lets us lift functions to context. In other words, you can lift functions of type (Boolean, Boolean) => Boolean (such as ||, && etc.) to (A => Boolean, A => Boolean) => (A => Boolean). (Here you can find a great explanation on this concept of lifting.)

However the data structure Set itself has an Applicative instance available, which will be favored over [B]Applicative[A => B] instance. To prevent that, we will have to explicitly tell the compiler to treat the given set as a function. We define a following enrichment for that:

scala> implicit def setAsFunction[A](set: Set[A]) = new {
     |   def f: A => Boolean = set
     | }
setAsFunction: [A](set: Set[A])java.lang.Object{def f: A => Boolean}

scala> Set(3, 4, 2).f
res144: Int => Boolean = Set(3, 4, 2)

And now put this Applicative goodness into use.

scala> val map = Map("1" -> 1, "2" -> 2, "3" -> 3)
map: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2, 3 -> 3)

scala> map filterKeys ((Set("1", "2").f |@| Set("2", "3").f)(_ && _))
res150: scala.collection.immutable.Map[java.lang.String,Int] = Map(2 -> 2)

scala> map filterKeys ((Set("1", "2").f |@| Set("2", "3").f)(_ || _))
res151: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2, 3 -> 3)

scala> map filterKeys (Set("2", "3").f map (!_))
res152: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1)

Note: All of the above requires Scalaz.

Upvotes: 9

oxbow_lakes
oxbow_lakes

Reputation: 134270

Answering the Question

You can take advantage of the fact that a Set[A] is a predicate; i.e. A => Boolean

map filterKeys set

Here it is at work:

scala> val map = Map("1" -> 1, "2" -> 2, "3" -> 3)
map: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2, 3 -> 3)

scala> val set = Set("1", "2")
set: scala.collection.immutable.Set[java.lang.String] = Set(1, 2)

scala> map filterKeys set
res0: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2)

Or if you prefer:

scala> map filterKeys Set("1", "2")
res1: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2)

Predicates

It's actually really useful to have some wrapper around a predicate. Like so:

scala> class PredicateW[A](self: A => Boolean) {
   | def and(other: A => Boolean): A => Boolean = a => self(a) && other(a)
   | def or(other: A => Boolean): A => Boolean = a => self(a) || other(a)
   | def unary_! : A => Boolean = a => !self(a)
   | }
defined class PredicateW

And an implicit conversion:

scala> implicit def Predicate_Is_PredicateW[A](p: A => Boolean) = new PredicateW(p)
Predicate_Is_PredicateW: [A](p: A => Boolean)PredicateW[A]

And then you can use it:

scala> map filterKeys (Set("1", "2") and Set("2", "3"))
res2: scala.collection.immutable.Map[java.lang.String,Int] = Map(2 -> 2)

scala> map filterKeys (Set("1", "2") or Set("2", "3"))
res3: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1, 2 -> 2, 3 -> 3)

scala> map filterKeys !Set("2", "3")
res4: scala.collection.immutable.Map[java.lang.String,Int] = Map(1 -> 1)

This can be extended to xor, nand etc etc and if you include symbolic unicode can make for amazingly readable code:

val mustReport = trades filter (uncoveredShort ∨ exceedsDollarMax)

val european = { 
  val Europe = (_ : Market).exchange.country.region == Region.EU
  trades filter (_.market ∈: Europe)
}

Upvotes: 54

Jesper
Jesper

Reputation: 206816

Sorry, not a direct answer to your question, but if you know which keys you want to remove (instead of which ones you want to keep), you could do this:

map -- Set("3")

Upvotes: 9

Related Questions