MultiSet: plus, minus, intersect, and improved equals

kotlin-utils/build.gradle.kts

@@ -2,11 +2,11 @@ plugins {
     `java-library`
     `maven-publish`
     id("org.jetbrains.kotlin.jvm") version "1.5.31"
-    id("org.jlleitschuh.gradle.ktlint") version "10.3.0" // ktlint
+    id("org.jlleitschuh.gradle.ktlint") version "11.0.0" // ktlint
 }
 
 group = "xyz.lbres"
-version = "0.3.1"
+version = "0.4.0"
 
 repositories {
     mavenCentral()

kotlin-utils/src/main/kotlin/xyz/lbres/kotlinutils/set/multiset/MultiSet.kt

@@ -6,5 +6,43 @@ package xyz.lbres.kotlinutils.set.multiset
  * Read/write access is available through the [MutableMultiSet] interface.
  */
 interface MultiSet<E> : Set<E> {
+    /**
+     * All distinct values contained in the MultiSet
+     */
+    val distinctValues: Set<E>
+
+    /**
+     * Get the number of occurrences of a given element.
+     *
+     * @param element [E]
+     * @return [Int]: The number of occurrences of [element]. 0 if the element does not exist.
+     */
     fun getCountOf(element: E): Int
+
+    /**
+     * Create a new MultiSet with values that are in this set but not the other set.
+     * If there are multiple occurrences of a value, the number of occurrences in the other set will be subtracted from the number in this MultiSet.
+     *
+     * @param other [MultiSet]<[E]>: values to subtract from this MultiSet
+     * @return [MultiSet]<[E]>: MultiSet containing the items in this MultiSet but not the other
+     */
+    operator fun minus(other: MultiSet<E>): MultiSet<E>
+
+    /**
+     * Create a new MultiSet with all values from both sets.
+     * If there are multiple occurrences of a value, the number of occurrences in the other set will be added to the number in this MultiSet.
+     *
+     * @param other [MultiSet]<[E]>: values to add to this MultiSet
+     * @return [MultiSet]<[E]>: MultiSet containing all values from both MultiSets
+     */
+    operator fun plus(other: MultiSet<E>): MultiSet<E>
+
+    /**
+     * Create a new MultiSet with values that are shared between the sets.
+     * If there are multiple occurrences of a value, the smaller number of occurrences will be used.
+     *
+     * @param other [MultiSet]<[E]>: values to intersect with this MultiSet
+     * @return [MultiSet]<[E]>: MultiSet containing only values that are in both MultiSets
+     */
+    fun intersect(other: MultiSet<E>): MultiSet<E>
 }

kotlin-utils/src/main/kotlin/xyz/lbres/kotlinutils/set/multiset/MultiSetImpl.kt

@@ -1,51 +1,76 @@
 package xyz.lbres.kotlinutils.set.multiset
 
+import kotlin.math.max
+import kotlin.math.min
+
 /**
  * Set implementation that allows multiple occurrences of the same value.
  */
-internal class MultiSetImpl<E> constructor(elements: Collection<E>) : MultiSet<E> {
+internal class MultiSetImpl<E> : MultiSet<E> {
     /**
      * Number of elements in set.
      */
-    override val size: Int = elements.size
+    override val size: Int
+
+    /**
+     * All distinct values contained in the MultiSet, without any counts
+     */
+    override val distinctValues: Set<E>
 
     /**
      * Store the number of occurrences of each element in set.
      * Counts are guaranteed to be greater than 0.
      */
-    private val countsMap: HashMap<E, Int>
+    private val countsMap: Map<E, Int>
 
     /**
      * The initial elements that were passed to the constructor.
      */
-    private val initialElements: Collection<E> = elements
+    private val initialElements: Collection<E>
 
     /**
      * String representation of the set.
      */
     private val string: String
 
     /**
-     * Initialize stored variables.
+     * Initialize stored variables from a collection of values.
      */
-    init {
-        // init string
-        string = if (size == 0) {
-            "[]"
-        } else {
-            val elementsString = elements.joinToString(", ")
-            "[$elementsString]"
-        }
+    constructor(elements: Collection<E>) {
+        size = elements.size
+        initialElements = elements
+        string = createString()
 
         // init counts
         val mutableMap: MutableMap<E, Int> = mutableMapOf()
+        val mutableValues: MutableSet<E> = mutableSetOf()
 
         for (element in elements) {
             val currentCount = mutableMap[element] ?: 0
             mutableMap[element] = currentCount + 1
+            mutableValues.add(element)
         }
 
-        countsMap = HashMap(mutableMap)
+        // cast to simpler data structures
+        countsMap = mutableMap.toMap()
+        distinctValues = mutableValues.toSet()
+    }
+
+    /**
+     * Initialize stored variables from an existing counts map.
+     */
+    private constructor(counts: Map<E, Int>) {
+        countsMap = counts
+        size = counts.values.fold(0, Int::plus)
+        distinctValues = counts.keys
+
+        initialElements = counts.flatMap {
+            val element = it.key
+            val count = it.value
+            List(count) { element }
+        }
+
+        string = createString()
     }
 
     /**
@@ -72,6 +97,59 @@ internal class MultiSetImpl<E> constructor(elements: Collection<E>) : MultiSet<E
         return newSet.countsMap.all { countsMap.contains(it.key) && it.value <= getCountOf(it.key) }
     }
 
+    /**
+     * Create a new MultiSet with values that are in this set but not the other set.
+     * If there are multiple occurrences of a value, the number of occurrences in the other set will be subtracted from the number in this MultiSet.
+     *
+     * @param other [MultiSet]<[E]>: MultiSet to subtract from current
+     * @return [MultiSet]<[E]>: MultiSet containing the items in this MultiSet but not the other
+     */
+    override operator fun minus(other: MultiSet<E>): MultiSet<E> {
+        val newCounts = countsMap.keys.associateWith {
+            val count = getCountOf(it)
+            val otherCount = other.getCountOf(it)
+            max(count - otherCount, 0)
+        }.filter { it.value > 0 }.toMap()
+
+        return MultiSetImpl(newCounts)
+    }
+
+    /**
+     * Create a new MultiSet with all values from both sets.
+     * If there are multiple occurrences of a value, the number of occurrences in the other set will be added to the number in this MultiSet.
+     *
+     * @param other [MultiSet]<[E]>: MultiSet to add to current
+     * @return [MultiSet]<[E]>: MultiSet containing all values from both MultiSets
+     */
+    override operator fun plus(other: MultiSet<E>): MultiSet<E> {
+        val allValues = distinctValues + other.distinctValues
+
+        val newCounts = allValues.associateWith {
+            getCountOf(it) + other.getCountOf(it)
+        }
+
+        return MultiSetImpl(newCounts)
+    }
+
+    /**
+     * Create a new MultiSet with values that are shared between the sets.
+     * If there are multiple occurrences of a value, the smaller number of occurrences will be used.
+     *
+     * @param other [MultiSet]<[E]>: MultiSet to intersect with current
+     * @return [MultiSet]<[E]>: MultiSet containing only values that are in both MultiSets
+     */
+    override fun intersect(other: MultiSet<E>): MultiSet<E> {
+        val allValues = distinctValues + other.distinctValues
+
+        val newCounts = allValues.associateWith {
+            val count = getCountOf(it)
+            val otherCount = other.getCountOf(it)
+            min(count, otherCount)
+        }.filter { it.value > 0 }.toMap()
+
+        return MultiSetImpl(newCounts)
+    }
+
     /**
      * If the current set contains 0 elements.
      *
@@ -80,7 +158,7 @@ internal class MultiSetImpl<E> constructor(elements: Collection<E>) : MultiSet<E
     override fun isEmpty(): Boolean = countsMap.isEmpty()
 
     /**
-     * Get the number of occurrences of a given element in the current set.
+     * Get the number of occurrences of a given element.
      *
      * @param element [E]
      * @return [Int]: the number of occurrences of [element]. 0 if the element does not exist.
@@ -94,11 +172,29 @@ internal class MultiSetImpl<E> constructor(elements: Collection<E>) : MultiSet<E
      * @return [Boolean]: true if [other] is a non-null ImmutableMultiSet which contains the same values as the current set, false otherwise
      */
     override fun equals(other: Any?): Boolean {
-        if (other == null || other !is MultiSetImpl<*>) {
+        if (other == null || other !is MultiSet<*>) {
             return false
         }
 
-        return countsMap == other.countsMap
+        return try {
+            other as MultiSet<E>
+            return minus(other).distinctValues.isEmpty() && other.minus(this).distinctValues.isEmpty()
+        } catch (e: Exception) {
+            false
+        }
+    }
+
+    /**
+     * Create the static string representation of the set.
+     * Stored in a helper so it can be reused in both constructors.
+     */
+    private fun createString(): String {
+        if (initialElements.isEmpty()) {
+            return "[]"
+        }
+
+        val elementsString = initialElements.joinToString(", ")
+        return "[$elementsString]"
     }
 
     /**
@@ -115,5 +211,5 @@ internal class MultiSetImpl<E> constructor(elements: Collection<E>) : MultiSet<E
      */
     override fun toString(): String = string
 
-    override fun hashCode(): Int = listOf("ImmutableMultiSet", countsMap).hashCode()
+    override fun hashCode(): Int = listOf(javaClass.name, countsMap).hashCode()
 }