Add groupBy for NonEmptySet

core/src/main/scala/cats/data/NonEmptySet.scala

@@ -346,6 +346,21 @@ sealed class NonEmptySetOps[A](val value: NonEmptySet[A]) {
     */
   def zipWithIndex: NonEmptySet[(A, Int)] =
     NonEmptySetImpl.create(toSortedSet.zipWithIndex)
+
+  /**
+    * Groups elements inside this `NonEmptySet` according to the `Order`
+    * of the keys produced by the given mapping function.
+    */
+  def groupBy[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NonEmptySet[A]] = {
+     reduceLeftTo(a => NonEmptyMap.one(f(a), NonEmptySet.one(a))) { (acc, a) =>
+       val key = f(a)
+       val result = acc.lookup(key) match {
+         case Some(nes) => nes.add(a)
+         case _ => NonEmptySet.one(a)
+       }
+       acc.add((key, result))
+     }
+  }
 }
 
 private[data] sealed abstract class NonEmptySetInstances {

core/src/main/scala/cats/syntax/all.scala

@@ -64,3 +64,4 @@ trait AllSyntaxBinCompat1
     with NestedSyntax
     with BinestedSyntax
     with ParallelFlatSyntax
+    with SetSyntax

core/src/main/scala/cats/syntax/set.scala

@@ -0,0 +1,52 @@
+package cats.syntax
+
+import scala.collection.immutable.{SortedSet, SortedMap}
+import cats.data.NonEmptySet
+import cats.Order
+
+trait SetSyntax {
+  implicit final def catsSyntaxSet[A](se: SortedSet[A]): SetOps[A] = new SetOps(se)
+}
+
+final class SetOps[A](val se: SortedSet[A]) extends AnyVal {
+
+  /**
+    * Returns an Option of NonEmptySet from a SortedSet
+    *
+    * Example:
+    * {{{
+    * scala> import scala.collection.immutable.SortedSet
+    * scala> import cats.data.NonEmptySet
+    * scala> import cats.implicits._
+    *
+    * scala> val result1: SortedSet[Int] = SortedSet(1, 2)
+    * scala> result1.toNes
+    * res0: Option[NonEmptySet[Int]] = Some(TreeSet(1, 2))
+    *
+    * scala> val result2: SortedSet[Int] = SortedSet.empty[Int]
+    * scala> result2.toNes
+    * res1: Option[NonEmptySet[Int]] = None
+    * }}}
+    */
+  def toNes: Option[NonEmptySet[A]] = NonEmptySet.fromSet(se)
+
+  /**
+    * Groups elements inside this `SortedSet` according to the `Order` of the keys
+    * produced by the given mapping function.
+    *
+    * {{{
+    * scala> import cats.data.NonEmptySet
+    * scala> import scala.collection.immutable.{SortedMap, SortedSet}
+    * scala> import cats.implicits._
+    *
+    * scala> val sortedSet = SortedSet(12, -2, 3, -5)
+    *
+    * scala> sortedSet.groupByNes(_ >= 0)
+    * res0: SortedMap[Boolean, NonEmptySet[Int]] = Map(false -> TreeSet(-5, -2), true -> TreeSet(3, 12))
+    * }}}
+    */
+  def groupByNes[B](f: A => B)(implicit B: Order[B]): SortedMap[B, NonEmptySet[A]] = {
+    implicit val ordering = B.toOrdering
+    toNes.fold(SortedMap.empty[B, NonEmptySet[A]])(_.groupBy(f).toSortedMap)
+  }
+}

tests/src/test/scala/cats/tests/NonEmptySetSuite.scala

@@ -233,4 +233,11 @@ class NonEmptySetSuite extends CatsSuite {
     }
   }
 
+  test("NonEmptySet#groupBy is consistent with Set#groupBy") {
+    forAll { (nes: NonEmptySet[Int], f: Int => Int) =>
+      nes.groupBy(f).map(_.toSortedSet).toSortedMap should === (nes.toSortedSet.groupBy(f))
+    }
+  }
+
+
 }