본문 바로가기
알고리즘 문제 풀이/DFS

938.Range Sum of BST

by 가나무마 2021. 5. 6.
728x90

주어진 트리 노드에서 노드의 값을 val이라고 정의했을 때,

low <= node.val <= high 조건을 만족하는 값들의 총합을 구하시오.

[처음 푼 코틀린 코드]

/**
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    fun rangeSumBST(root: TreeNode?, low: Int, high: Int): Int {
        var stack = Stack<TreeNode>()
        stack.push(root)
        //총합을 담을 변수.
        var sum : Int = 0

        while (!stack.empty()) {
            var node = stack.pop()
            sum += checkRange(stack,node,low,high)
        }

        return sum
    }
    fun checkRange(stack:Stack<TreeNode>, node:TreeNode, low:Int,high:Int) : Int {
        //node의 값이 low와 high사이에 있으면 노드의 좌우를 null 체크하고 stack에 넣어줌.
        if (low <= node.`val` && node.`val` <= high) {
            if (node.left != null) {
                stack.push(node.left)
            }
            if (node.right != null) {
                stack.push(node.right)
            }
            return node.`val`
        } 
        //node의 값이 low보다 작으면 left의 값이 node의 값보다 작기 때문에 
        //무조건 성립하지 않음.
        //따라서 node의 값보다 큰 right만 stack에 넣어줌.
        else if (node.`val` < low) {
            if (node.right != null) {
                stack.push(node.right)
            }
        } 
        //node의 값이 high보다 크면 right의 값이 node의 값보다 크기 때문에 
        //무조건 성립하지 않음.
        //따라서 node의 값보다 작은 left만 stack에 넣어줌.
        else if (high < node.`val` ) {
            if (node.left != null) {
                stack.push(node.left)
            }
        }
        return 0
    }
}

Null 처리하고 stack에 넣는 게 중복돼서 메소드로 분리했음.

/**
 * Example:
 * var ti = TreeNode(5)
 * var v = ti.`val`
 * Definition for a binary tree node.
 * class TreeNode(var `val`: Int) {
 *     var left: TreeNode? = null
 *     var right: TreeNode? = null
 * }
 */
class Solution {
    fun rangeSumBST(root: TreeNode?, low: Int, high: Int): Int {
        var stack = Stack<TreeNode>()
        stack.push(root)
        var sum : Int = 0

        while (!stack.empty()) {
            var node = stack.pop()
            sum += checkRange(stack,node,low,high)
        }

        return sum
    }
    fun checkRange(stack:Stack<TreeNode>, node:TreeNode, low:Int,high:Int) : Int {
        if (low <= node.`val` && node.`val` <= high) {
            nullCheckAndPush(stack, node.left)
            nullCheckAndPush(stack, node.right)

            return node.`val`
        } else if (node.`val` < low) {
            nullCheckAndPush(stack, node.right)
        } else if (high < node.`val` ) {
            nullCheckAndPush(stack, node.left)
        }
        return 0
    }
    //node가 null이 아니면 stack에 넣어줌.
    fun nullCheckAndPush(stack:Stack<TreeNode>, node:TreeNode?) {
        if (node != null) {
            stack.push(node)
        }
    }
}
728x90
반응형