🙆

ABC262 C - Min Max Pairのメモ

2022/12/27に公開

目についたC問題を解いてみるの回です。比較的最近のだけど参加できなかった回のやつです。

https://atcoder.jp/contests/abc262/tasks/abc262_c

考察

1-indexで考えた場合は
a_i = iかつa_j = j
a_i = jかつa_j = i

のどちらかの場合に条件を満たします。入力例1でいうと、(1,4)が①、(2,3)が②です。

これは、それぞれ別の処理として考えるとしっくりきました。

①のケース

これは、iとjの両方について「数列の何番目か」と「その位置の値」が一致する場合ということです。

入力例2で考えるとわかりやすかったです。

[5, 8, 2, 2, 1, 6, 7, 2, 9, 10]

この場合、[6, 7, 9, 10]がその条件を満たします。これらの組み合わせを数えると、①の分の答えになります。
全部列挙すると(6, 7), (6, 9), (6, 10), (7, 9), (7, 10), (9, 10)の6通りです。

これを愚直に列挙すると間に合わないですが、欲しいのは「何通りか」という数値だけなので計算によって定数時間で求められます。
これは一般的な組み合わせの計算で求められます。(覚えてなかったのでググりましたが…)
_nC_rです。これはn - (n - 1) ÷ rで求められます。
n個のうちr個選んだ組み合わせという意味なので、上記の場合は_4C_2です。なので4 - (4 - 1) ÷ 2で6と求められました。

この計算で必要になったのは4と2という数字だけです。4は[6, 7, 9, 10]の個数ですが、個数さえあれば具体的にどの数なのかは関係ないことになります。2は固定で2です。

なので、①に関しては「数列の何番目か」と「その位置の値」が一致している要素が何個あるかを数えて、上記の計算をすれば求められます。

②のケース

こちらはもっと簡単で、単純に数えればいいです。

[5, 8, 2, 2, 1, 6, 7, 2, 9, 10]

たとえば数列の1番目の要素は5です。また、数列の5番目の要素は1です。このように「数列の何番目か」と「その位置の値」が入れ替わっているようなものを数えればいいです。
他には数列の2番目は8、数列の8番目は2、という関係があります。それ以外にはないので、②の場合の答えは2です。

①で求めた6と②で求めた2を足せば8になります。これは出力例2と一致します。

実装

import kotlin.math.max
import kotlin.math.min

fun main() {
    val n = readLine()!!.toInt()
    val a = listOf(-1) + readLine()!!.split(" ").map { it.toInt() }

    // ①順番通りになっているものを数える
    val orderedCount = a.filterIndexed { index, i -> index == i }.count().toLong()

    var ans = orderedCount * (orderedCount - 1) / 2

    // ②入れ替わっているものを数える
    for(i in 1 .. n) {
        val j = a[i]
        if(i < j && a[j] == i) {
            ans++
        }
    }

    println(ans)
}

1-indexedに補正するために先頭に-1を入れています。ここに0を入れるとそれもカウントされてしまうので0以外の値にします。

filterIndexed を使って「数列の何番目か」と「その位置の値」が一致する要素をカウントしています。注意点として、最大ケースの場合は500000と499999を掛けることになるのでIntを範囲を超えてしまいます。なのでLongにします。

蛇足

最初は順番が入れ替わっている場合があるのが嫌で、ソートできないかなと考えてそれで提出してWAとか出していました…

実際にはソートする必要はなく、というかソートしたらダメでした。

Discussion