Merge pull request #4711 from BartekBH/main

Optimize distinctBy implementation for non-empty collections
typelevel · Feb 1, 2025 · 5ac7e64 · 5ac7e64
2 parents f5549aa + 19f2d92
commit 5ac7e64
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 30 deletions.
diff --git a/core/src/main/scala/cats/data/NonEmptyList.scala b/core/src/main/scala/cats/data/NonEmptyList.scala
@@ -343,18 +343,24 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) extends NonEmptyCollec
   override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyList[AA] = distinctBy(identity[AA])
 
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyList[A] = {
-    implicit val ord: Ordering[B] = O.toOrdering
-
-    val buf = ListBuffer.empty[A]
-    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
-      val b = f(a)
-      if (elementsSoFar(b)) elementsSoFar
-      else {
-        buf += a; elementsSoFar + b
+    if (tail.isEmpty) this
+    else {
+      implicit val ord: Ordering[B] = O.toOrdering
+
+      val bldr = List.newBuilder[A]
+      val seen = mutable.TreeSet.empty[B]
+      var rest = tail
+      seen.add(f(head))
+      while (rest.nonEmpty) {
+        val next = rest.head
+        rest = rest.tail
+        if (seen.add(f(next))) {
+          bldr += next
+        }
       }
-    }
 
-    NonEmptyList(head, buf.toList)
+      NonEmptyList(head, bldr.result())
+    }
   }
 
   /**

diff --git a/core/src/main/scala/cats/data/NonEmptySeq.scala b/core/src/main/scala/cats/data/NonEmptySeq.scala
@@ -239,18 +239,21 @@ final class NonEmptySeq[+A] private (val toSeq: Seq[A]) extends AnyVal with NonE
   override def distinct[AA >: A](implicit O: Order[AA]): NonEmptySeq[AA] = distinctBy(identity[AA])
 
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptySeq[A] = {
-    implicit val ord: Ordering[B] = O.toOrdering
-
-    val buf = Seq.newBuilder[A]
-    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
-      val b = f(a)
-      if (elementsSoFar(b)) elementsSoFar
-      else {
-        buf += a; elementsSoFar + b
+    if (toSeq.lengthCompare(1) == 0) this
+    else {
+      implicit val ord: Ordering[B] = O.toOrdering
+
+      val bldr = Seq.newBuilder[A]
+      val seen = mutable.TreeSet.empty[B]
+      val it = iterator
+      while (it.hasNext) {
+        val next = it.next()
+        if (seen.add(f(next)))
+          bldr += next
       }
-    }
 
-    NonEmptySeq(head, buf.result())
+      NonEmptySeq.fromSeqUnsafe(bldr.result())
+    }
   }
 
   /**

diff --git a/core/src/main/scala/cats/data/NonEmptyVector.scala b/core/src/main/scala/cats/data/NonEmptyVector.scala
@@ -249,18 +249,21 @@ final class NonEmptyVector[+A] private (val toVector: Vector[A])
   override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyVector[AA] = distinctBy(identity[AA])
 
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyVector[A] = {
-    implicit val ord: Ordering[B] = O.toOrdering
-
-    val buf = Vector.newBuilder[A]
-    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
-      val b = f(a)
-      if (elementsSoFar(b)) elementsSoFar
-      else {
-        buf += a; elementsSoFar + b
+    if (toVector.lengthCompare(1) == 0) this
+    else {
+      implicit val ord: Ordering[B] = O.toOrdering
+
+      val bldr = Vector.newBuilder[A]
+      val seen = mutable.TreeSet.empty[B]
+      val it = iterator
+      while (it.hasNext) {
+        val next = it.next()
+        if (seen.add(f(next)))
+          bldr += next
       }
-    }
 
-    NonEmptyVector(head, buf.result())
+      NonEmptyVector.fromVectorUnsafe(bldr.result())
+    }
   }
 
   /**