AtCoder Beginner Contest 170 E - Smart Infants

実装がかなり重かった上に定数倍がきつそうな感じがしたけど、よく見ると制限3.5secだったのでギリギリ間に合った。

実装

こういう問題はまず与えられた情報をとりあえず書いてみると思い浮かんだりする。とりあえずこの2つはすぐわかる。

幼児からレートを検索するための配列rating
幼児から所属幼稚園を検索するための配列belongs

val belongs = IntArray(N + 1)
val rating = IntArray(N + 1)

次に、各クエリにおいて

最高レートが更新される可能性のある園は最高2つまで
毎回最小値を求める必要がある

これは一点更新区間取得なので、それを処理するためのSegmentTreeを用意する。演算はmin()、単位元はInt.MAX_VALUE。 ratingInKinderの説明はあとで。

val st = IntSegTree(
    ratingInKinder.map { if (it.isEmpty()) Int.MAX_VALUE else it.lastKey() }.toIntArray(),
    { i1, i2 -> min(i1, i2) },
    Int.MAX_VALUE
)

ここからこのSegmentTreeを更新するために必要な情報はなにかを考える。最高レートが変化する場合や、園から一人もいなくなる場合など、色々な場合を考える。すると、以下のような情報があれば更新できると思いつく。

その園に存在するレーティングと、その人数
その園に存在するレーティングの降順のSet

これは、いわゆる順序付き多重集合というもので保持できる。C++にはこのためのmultisetがあるらしいが、kotlinには無いためTreeMap(keyが自動ソートされたMap)を使って、valueとして人数を保持する。

//初期化
val ratingInKinder = Array(2 * 100000 + 1) { TreeMap<Int, Int>() }
//Ai Biの情報から最初に更新する処理
ratingInKinder[B][A] = (ratingInKinder[B][A] ?: 0) + 1

持っていないといけない情報が多いが、これでクエリを処理することができる。実際に書いたコードでは、わかりやすさのために前の所属先と次の所属先の最高レートの情報を毎回更新している。
そのレートの幼児がいなくなる場合にremoveすることや、園に一人もいなくなる場合にセグ木をMAXで更新する処理を忘れないように注意。

val child = // Ci
val before = belongs[child]
val after = // Di

//前の所属先に関する処理
if (ratingInKinder[before][rating[child]]!! == 1) {
    ratingInKinder[before].remove(rating[child])
} else {
    ratingInKinder[before][rating[child]] = ratingInKinder[before][rating[child]]!! - 1
}
if (ratingInKinder[before].isEmpty()) {
    st.update(before, Int.MAX_VALUE)
} else {
    st.update(before, ratingInKinder[before].lastKey())
}

//次の所属先に関する処理
if (!ratingInKinder[after].containsKey(rating[child])) {
    ratingInKinder[after].put(rating[child], 1)
} else {
    ratingInKinder[after][rating[child]] = ratingInKinder[after][rating[child]]!! + 1
}
st.update(after, ratingInKinder[after].lastKey())

belongs[child] = after
ans[i] = st.getOp(1, ratingInKinder.size)

最後に答えを出力。java/kotlinはprintlnが遅いので一回で出力したほうがいい。

println(ans.joinToString(separator = "\n"))

ソースコード