50 lines
1.5 KiB
Kotlin
50 lines
1.5 KiB
Kotlin
package moe.odango.index.utils
|
|
|
|
class MergeMap<T> {
|
|
private var bucketIndex = 0L
|
|
private val buckets: MutableMap<Long, MutableSet<T>> = mutableMapOf()
|
|
private val bucketIndexByItem: MutableMap<T, Long> = mutableMapOf()
|
|
|
|
fun add(a: T, b: T) {
|
|
val bucketA = bucketIndexByItem[a]
|
|
val bucketB = bucketIndexByItem[b]
|
|
if (bucketA == null && bucketB == null) {
|
|
bucketIndexByItem[a] = bucketIndex
|
|
bucketIndexByItem[b] = bucketIndex
|
|
buckets[bucketIndex] = mutableSetOf(a, b)
|
|
bucketIndex++
|
|
return
|
|
}
|
|
|
|
if (bucketB == bucketA) {
|
|
return
|
|
}
|
|
|
|
if (bucketA == null && bucketB != null) {
|
|
bucketIndexByItem[a] = bucketB
|
|
buckets.getOrPut(bucketB, ::mutableSetOf).add(a)
|
|
return
|
|
}
|
|
|
|
if (bucketB == null && bucketA != null) {
|
|
bucketIndexByItem[b] = bucketA
|
|
buckets.getOrPut(bucketA, ::mutableSetOf).add(b)
|
|
return
|
|
}
|
|
|
|
// Always false :)
|
|
if (bucketA == null || bucketB == null) return
|
|
|
|
val bucket = buckets.getOrPut(bucketB, ::mutableSetOf)
|
|
buckets.getOrPut(bucketA, ::mutableSetOf).addAll(bucket)
|
|
for (item in bucket) {
|
|
bucketIndexByItem[item] = bucketA
|
|
}
|
|
}
|
|
|
|
operator fun get(item: T): Set<T>? {
|
|
return buckets[bucketIndexByItem[item] ?: return null]
|
|
}
|
|
|
|
operator fun iterator(): Iterator<Set<T>> = buckets.values.iterator()
|
|
}
|