알고리즘

2024.10.23 2641. Cousins in Binary Tree II

물빠진떡 2024. 10. 23. 17:40

문제

Given the root of a binary tree, replace the value of each node in the tree with the sum of all its cousins' values.

Two nodes of a binary tree are cousins if they have the same depth with different parents.

Return the root of the modified tree.

Note that the depth of a node is the number of edges in the path from the root node to it.

요약

트리 문제이다. 노드수의 제약이 10000개 밖에 되지않아 충분히 전탐으로 될꺼라고 생각했다.

풀이

한번 돌면서 한 depth의 총합을 구하고

다시 한번 돌면서 현재 노드의 두 자식 val을 (아래 depth의 총합) - (두 자식의 총합)으로 해주면된다.

그리고 root는 부모가 없으므로 무조건 0가 된다.

코드

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public TreeNode replaceValueInTree(TreeNode root) {
        int[] totals = new int[100_001];
        Queue<DepthNode> queue = new LinkedList<>();
        queue.add(new DepthNode(root, 1));
        while (!queue.isEmpty()) {
            DepthNode node = queue.poll();
            totals[node.depth] += node.val;
            if (node.left != null) {
                queue.add(new DepthNode(node.left, node.depth + 1));
            }
            if (node.right != null) {
                queue.add(new DepthNode(node.right, node.depth + 1));
            }
        }
        dfs(totals, 1, root);
        root.val = 0;
        return root;
    }

    private void dfs(int[] totals, int depth, TreeNode node) {
        int sum = 0;
        if (node.left != null) {
            dfs(totals, depth + 1, node.left);
            sum += node.left.val;
        }

        if (node.right != null) {
            dfs(totals, depth + 1, node.right);
            sum += node.right.val;
            node.right.val = totals[depth + 1] - sum;
        }

        if (node.left != null) {
            node.left.val = totals[depth + 1] - sum;
        }
    }

    private class DepthNode{
        TreeNode node;
        int depth;
        TreeNode left;
        TreeNode right;
        int val;

        DepthNode(TreeNode node, int depth) {
            this.node = node;
            this.depth = depth;
            this.right = node.right;
            this.left = node.left;
            this.val = node.val;
        }
    }
}

풀이 시간

10분

체감 난이도

easy