Last active
June 14, 2021 14:18
-
-
Save malterb/ac557e5c77471c1ead8d6274094dc5a0 to your computer and use it in GitHub Desktop.
groupBy on scala iterator
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package utils | |
import scala.collection.mutable | |
import scala.language.implicitConversions | |
object IteratorUtils { | |
/** | |
* Implicit class that adds "groupBy" functionality to an iterator. The result is a mutable.HashMap and will be | |
* stored in memory in its entirety | |
* @param iterator the iterator to operate on | |
* @tparam ElementType type of iterator elements | |
*/ | |
implicit class IteratorGroupBy[ElementType](iterator: Iterator[ElementType]) { | |
/** | |
* groups each element in iterator by the key function and stores them in a vector | |
* @param makeKey the function to get the key from an element | |
* @tparam KeyType self expl. | |
* @example | |
* case class Toy(id: Int, price: Int) | |
* val itCollect = Iterator(Toy(1,100),Toy(1,200),Toy(2,200),Toy(3,10),Toy(1,2234)) | |
* itCollect.groupByCollect(_.id) | |
* _______________________________________ | |
* res0: scala.collection.mutable.HashMap[Int,Vector[Toy]] = Map(2 -> Vector(Toy(2,200)), 1 -> Vector(Toy(1,2234), Toy(1,200), Toy(1,100)), 3 -> Vector(Toy(3,10))) | |
*/ | |
def groupByCollect[KeyType](makeKey: (ElementType) => KeyType): mutable.HashMap[KeyType, Vector[ElementType]] = { | |
val internal = new mutable.HashMap[KeyType, Vector[ElementType]]() | |
iterator.foreach(elementA => { | |
val key = makeKey(elementA) | |
val res = internal.get(key) | |
internal(key) = res.map(value => elementA +: value).getOrElse(Vector(elementA)) | |
}) | |
internal | |
} | |
/** | |
* Keeps only one element per key from iterator. compare function is used to merge/compare/fold two different elements with the same key | |
* @param makeKey the function to get the key from an element | |
* @param compare given two elements gives you the element to store | |
* @tparam KeyType self expl. | |
* @example | |
* case class Toy(id: Int, price: Int) | |
* val itFold = Iterator(Toy(1,100),Toy(1,200),Toy(2,200),Toy(3,10),Toy(1,2234)) | |
* itFold.groupByFold(_.id,(a,b) => if(a.price < b.price) a else b) | |
* _______________________________________ | |
* res0: scala.collection.mutable.HashMap[Int,Toy] = Map(2 -> Toy(2,200), 1 -> Toy(1,100), 3 -> Toy(3,10)) | |
*/ | |
def groupByFold[KeyType](makeKey: (ElementType) => KeyType, compare: (ElementType,ElementType) => ElementType): mutable.HashMap[KeyType, ElementType] = { | |
val internal = new mutable.HashMap[KeyType, ElementType]() | |
iterator.foreach(elementA => { | |
val key = makeKey(elementA) | |
val res = internal.get(key) | |
internal(key) = res.map(value => compare(elementA,value)).getOrElse(elementA) | |
}) | |
internal | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment