Mapping and heads for ArrayMutableSortedMap.

appliedtopology · Feb 6, 2024 · 3c94e19 · 3c94e19
1 parent 3ec6b91
commit 3c94e19
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 1 deletion.
diff --git a/src/commonMain/kotlin/org/appliedtopology/tda4j/Sets.kt b/src/commonMain/kotlin/org/appliedtopology/tda4j/Sets.kt
@@ -61,7 +61,7 @@ class ArrayMutableSortedSet<T : Comparable<T>>(capacity: Int = 8) : ArrayMutable
 }
 typealias MutableSortedSet<V> = ArrayMutableSortedSet<V>
 
-open class ArrayMutableSortedMap<K : Comparable<K>, V>(capacity: Int = 8, defaultValue: V? = null) :
+open class ArrayMutableSortedMap<K : Comparable<K>, V>(capacity: Int = 8, val defaultValue: V? = null) :
     ArrayMutableSortedSetBase<K>(capacity), MutableMap<K, V> {
     protected val _values: ArrayList<V> = ArrayList(capacity)
 
@@ -93,6 +93,32 @@ open class ArrayMutableSortedMap<K : Comparable<K>, V>(capacity: Int = 8, defaul
     override val values: MutableCollection<V>
         get() = _values.toMutableList()
 
+    fun ordinalKey(index: Int): K? = _set.getOrNull(index)
+
+    fun ordinalValue(index: Int): V? = _values.getOrElse(index) { i -> defaultValue }
+
+    fun ordinalItem(index: Int): MutableMap.MutableEntry<K, V>? {
+        val ok = ordinalKey(index)
+        val ov = ordinalValue(index)
+        if (ok == null || ov == null) {
+            return null
+        } else {
+            return PairEntry(ordinalKey(index)!!, ordinalValue(index)!!)
+        }
+    }
+
+    val headKey: K?
+        get() = ordinalKey(0)
+    val headValue: V?
+        get() = ordinalValue(0)
+    val headItem: MutableMap.MutableEntry<K, V>?
+        get() = ordinalItem(0)
+
+    fun replaceAllValues(transform: (V) -> V) {
+        for (idx in _values.indices)
+            _values[idx] = transform(_values[idx])
+    }
+
     override fun remove(key: K): V? {
         val idx = _set.binarySearch(key)
         if (idx >= 0) {

diff --git a/src/commonTest/kotlin/org/appliedtopology/tda4j/SetSpec.kt b/src/commonTest/kotlin/org/appliedtopology/tda4j/SetSpec.kt
@@ -6,6 +6,7 @@ import io.kotest.matchers.booleans.shouldBeFalse
 import io.kotest.matchers.booleans.shouldBeTrue
 import io.kotest.matchers.collections.shouldBeSorted
 import io.kotest.matchers.collections.shouldContainExactly
+import io.kotest.matchers.equals.shouldBeEqual
 import io.kotest.property.Arb
 import io.kotest.property.arbitrary.int
 import io.kotest.property.arbitrary.list
@@ -70,3 +71,34 @@ class SetSpec : StringSpec({
         }
     }
 })
+
+class MapSpec : StringSpec({
+    "Mutable Map can take keys and values" {
+        checkAll<List<Pair<Int, Int>>> {
+            val keys = it.map { (k, v) -> k }
+            val vals = it.map { (k, v) -> v }
+            val mmap = ArrayMutableSortedMap<Int, Int>(capacity = it.size, defaultValue = 0)
+            it.forEach { (k, v) -> mmap.put(k, v) }
+
+            val hv = mmap.headValue ?: 0
+
+            withClue("mmap has the right number of entries") {
+                mmap.size.shouldBeEqual(keys.toSet().size)
+            }
+            withClue("mmap has the right head key") {
+                assume(it.isNotEmpty())
+                mmap.headKey?.shouldBeEqual(keys.min())
+            }
+            withClue("mmap has a mapValues") {
+                assume(it.isNotEmpty())
+                val mmap2 = mmap.mapValues { (k, v) -> v * v }
+                mmap2[keys.min()] == hv * hv
+            }
+            withClue("mmap has a replaceAllValues to map values in place") {
+                assume(it.isNotEmpty())
+                mmap.replaceAllValues { v -> v * v }
+                (mmap.headValue ?: 0).shouldBeEqual(hv * hv)
+            }
+        }
+    }
+})