Submission #71703652

---

Source Code Expand

fun main() {
    val cin = System.`in`.bufferedReader()
    val cout = System.out.bufferedWriter()

    class DSU(n: Int) {
        private val p = IntArray(n)
        var cnt = n

        init {
            for (i in 0 until n) p[i] = i
        }

        fun p(i: Int): Int {
            if (p[i] != i) p[i] = p(p[i])
            return p[i]
        }

        fun j(i: Int, j: Int) {
            val pi = p(i)
            val pj = p(j)
            if (pi != pj){
                p[i] = pj
                cnt--
            }
        }
    }

    var n = cin.readLine().toInt()
    var p = cin.readLine().split(" ").map { it.toInt() - 1 }
    val dsu = DSU(n)
    for (i in 0 until n) {
        dsu.j(i, p[i])
    }
    var ans = 0L
    val cnt= IntArray(n)
    for (i in 0 until n) {
        cnt[dsu.p(i)]++
    }
    for (v in cnt) if (v > 1) ans += v.toLong()*(v-1)/2

    cout.write("${ans}")
    cout.newLine()
    cout.flush()
}

Submission Info

Judge Result