LeetCode/kt/find-elements-in-a-contaminated-binary-tree.kt

class FindElements(root: TreeNode?) {
    private val seen: MutableSet<Int> = mutableSetOf()

    private fun bfs(root: TreeNode?) {
        val q = ArrayDeque<TreeNode?>()
        root?.`val` = 0
        q.add(root)

        while (q.isNotEmpty()) {
            val node = q.removeFirst()
            seen.add(node!!.`val`)

            listOf(node?.left, node?.right)
                .withIndex()
                .filter { it.value != null }
                .forEach { it ->
                    it.value?.`val` = node!!.`val` * 2 + 1 + it.index
                    q.add(it.value)
                }
        }
    }

    init {
        bfs(root)
    }

    fun find(target: Int): Boolean = seen.contains(target)
}