cs61b sp21 lec14 Disjoint Sets

impl.kt

interface DisjointSets {
    fun connect(p: Int, q: Int)
    fun isConnected(p: Int, q: Int): Boolean
}

class DisjointSetsByArray(val elements: Int) : DisjointSets {
    private val map = IntArray(elements) { -1 }

    private fun top(id: Int): Int {
        if (id !in map.indices) return -1
        var i = id
        while (map[i] >= 0) i = map[i]
        return i
    }
    private fun retop(id: Int): Int {
        val top = top(id)
        if (top < 0) return -1
        var i = id
        while (map[i] >= 0) i = map[i].also { map[i] = top }
        return top
    }
    private fun weight(i: Int, j: Int) =
        if (map[i] < map[j]) i to j else j to i

    override fun connect(p: Int, q: Int) {
        if (p == q || p !in map.indices || q !in map.indices) return
        val (i, j) = weight(top(p), top(q))
        map[i] += map[j]
        map[j] = i
    }
    override fun isConnected(p: Int, q: Int) =
        retop(p).let { it >= 0 && it == retop(q) }
}

class DisjointSetsByMapEtTree(val elements: Int) : DisjointSets {
    fun rangeOK(p: Int, q: Int) = p in 0 until elements && q in 0 until elements

    val map = HashMap<Int, Node>(elements)
    fun top(id: Int) = map.getOrPut(id, { Node(id) }).top

    override fun connect(p: Int, q: Int) {
        if (!rangeOK(p, q)) return
        val a = top(p)
        val b = top(q)
        if (a != b) b.prev = a
    }
    override fun isConnected(p: Int, q: Int) =
        rangeOK(p, q) && top(p) == top(q)
}

class Node(val id: Int, var prev: Node? = null) {
    // INVARIANT: top.prev == null
    val top: Node
        get() {
            var curr = this
            // while (curr.prev != null) curr = curr.prev!!
            while (true) {
                val p = curr.prev
                if (p == null) break
                curr = p
            }
            return curr
        }
}

test.kt

import kotlin.test.*

class Suit {
    // fun makeDJSets(elements: Int): DisjointSets = DisjointSetsByMapEtTree(elements)
    fun makeDJSets(elements: Int): DisjointSets = DisjointSetsByArray(elements)

    @Test fun `none any two connected initially`() {
        makeDJSets(2).apply {
            assertFalse(isConnected(0, 1))
            assertFalse(isConnected(1, 0))
        }
    }

    @Test fun `p connected to itself`() {
        makeDJSets(2).apply {
            assertTrue(isConnected(0, 0))
            assertTrue(isConnected(1, 1))
        }
    }

    @Test fun `connect p to itself should not crash`() {
        makeDJSets(2).apply {
            connect(0, 0)
            connect(1, 1)
            assertTrue(isConnected(0, 0))
            assertTrue(isConnected(1, 1))
        }
    }

    @Test fun `p exceeds size connected to none`() {
        makeDJSets(2).apply {
            assertFalse(isConnected(22, 0))
            assertFalse(isConnected(22, 1))
            connect(22, 0)
            connect(22, 1)
            assertFalse(isConnected(22, 0))
            assertFalse(isConnected(22, 1))
        }
    }

    @Test fun `connected elements connected`() {
        makeDJSets(2).apply {
            assertFalse(isConnected(0, 1))
            connect(0, 1)
            assertTrue(isConnected(0, 1))
        }
    }
}

• add test connect(p, p)
• enable DisjointSetsByArray output DOT graph