leetcode - frog position after t seconds(kotlin)

level - hard

자세한 설명과 내용은 아래 링크를 참고해주세요.
https://velog.io/@mdok1112/leetcode-frog-position-after-t-secondsjava

위의 링크에서는 설명과 자바코드가 들어있으며
여기서는 코틀린으로 작성한 코드만 있습니다.

class Solution {
    fun generateTrees(edges:Array<IntArray>, start:Int, trees:HashMap<Int, HashSet<Int>>): Unit {
        for(edge in edges) {
            if(edge[0] == start) {
                if(addTree(start, edge[1], trees)) {
                    generateTrees(edges, edge[1], trees);
                }
            } else if(edge[1] == start) {
                if(addTree(start, edge[0], trees)) {
                    generateTrees(edges, edge[0], trees);
                }
            }
        }
    }

    fun addTree(key:Int, value:Int, trees: HashMap<Int, HashSet<Int>>): Boolean {
        if(trees.containsKey(key) && !trees.containsKey(value)) {
            trees[key]!!.add(value);
            return true;
        } else if(!trees.containsKey(value)) {
            val tree = HashSet<Int>();
            tree.add(value);
            trees[key] = tree;
            return true;
        }

        return false;
    }

    fun frogPosition(n: Int, edges: Array<IntArray>, t: Int, target: Int): Double {
        val trees = HashMap<Int, HashSet<Int>>();
        generateTrees(edges, 1, trees);
        
        val counts = ArrayList<Int>();
        var buf = target;
        while(buf != 1) {
            for(key in trees.keys) {
                val tree = trees[key];
                if(tree!!.contains(buf)) {
                    buf = key;
                    counts.add(tree.size);
                    break;
                }
            }
        }
        
        var result = 0.0;
        if(counts.size <= t) {
            if(counts.size < t && trees.containsKey(target)) {
                return 0.0;
            }
            if(counts.size != 0 || buf == 1) {
                result = 1.0;
            }
            for(count in counts) {
                result /= count;
            }
        }
        
        return result;
    }
}