Dive into Ofuton

お布団に飛び込もう

KotlinでMultiSetを雑に実装する

AtCoderの問題を解いているときにMultiSetがあると助かる場面があったためざっくり実装。

import java.util.TreeMap
import java.util.TreeSet

class  SortedMultiSet<T>(
    comparator: Comparator<T>? = null,
    val map: TreeMap<T, Int> = TreeMap<T, Int>(comparator)
) :MutableSet<T> by map.keys{
    override fun add(element: T): Boolean {
        map[element] = (map[element]?: 0) + 1
        return true
    }
    override fun remove(element: T): Boolean {
        val x = map[element]
        if(x == null){
            return false
        }
        if(x == 1){
            map.remove(element)
        } else {
            map[element] = x - 1
        }
        return true
    }
}
  • できること
    • 同じ値を複数追加
  • できないこと
    • 同じ値を含む全要素ループ←同じ値は1回しかループされない