[LeetCode] 814. Binary Tree Pruning

https://leetcode.com/problems/binary-tree-pruning/

Binary Tree Pruning - LeetCode

Binary Tree Pruning.

 

 

값이 0과 1중 하나로 되어 있는 이진 트리가 주어졌을 때, 자식 트리에 1이 없는 노드는 삭제해야 하는 문제이다.

풀이법도 금세 생각이 나고 짜는데도 문제가 없었는데 효율성면에서 이게 맞나..? 하는 생각이 들었다.

노드 최대 개수가 200개라 시간복잡도는 별 상관이 없겠지만.

 

생각한 풀이는 이렇다.

 

1. 각 노드의 값을 바로 아래 자식노드들의 값의 합 + 해당 노드의 값으로 나타낸다. 부분합을 구할때와 비슷하게 노드 값의 합을 이용해 재구성하는 것이다. 이게 나타내는 것은 해당 노드부터 모든 자식노드들을 훑었을때 1의 개수가 된다.

 

위의 예제에서, 

Example1. 

1
		0
		0	1

↓

2
		1
		0	1

 

Example2.

1
0		1
0	0	0	1

↓

3
0		2
		0	1

 

Example3.

1
1				0
1		1		0		1
0

↓

5
3				1
1		1		0		1
0

int sum(TreeNode** node) {
	if (*node!=NULL) {
		(*node)->val += sum(&((*node)->left)) + sum(&((*node)->right));
		return (*node)->val;
	}
	return 0;
}

2. 간단하다. 바꾼 트리가 나타낸 것이 해당 노드부터 모든 자식노드들까지 1의 개수이므로 해당 노드의 값이 0이면 제거해주면 된다.

 

Example3.

5
3				1
1		1		0		1
0

↓

5
3				1
1		1				1

void chk(TreeNode** node) {
	if (*node==NULL) return;
	if ((*node)->val==0) {*node = NULL; return;}
	chk(&((*node)->left));
	chk(&((*node)->right));
}

3. 마지막으로 모든 노드를 훑으며 자식 노드들의 합이 해당 노드의 값이라면, 해당 노드의 값은 원래 0이었단 소리이므로 0으로, 아니라면 1로 돌려준다.

 

Example3.

5
3				1
1		1				1

↓

1
1				1
1		1				1

 

void reset(TreeNode** node) {
	int s = 0;
	if (*node==NULL) return;
	if ((*node)->left!=NULL) s+=(*node)->left->val;
	if ((*node)->right!=NULL) s+=(*node)->right->val;
	(*node)->val = s==(*node)->val? 0:1;
	reset(&((*node)->left));
	reset(&((*node)->right));
}

 

전체 코드는 아래와 같다.

 

class Solution {
public:
    int sum(TreeNode** node) {
        if (*node!=NULL) {
            (*node)->val += sum(&((*node)->left)) + sum(&((*node)->right));
            return (*node)->val;
        }
        return 0;
    }
    void chk(TreeNode** node) {
        if (*node==NULL) return;
        if ((*node)->val==0) {*node = NULL; return;}
        chk(&((*node)->left));
        chk(&((*node)->right));
    }
    void reset(TreeNode** node) {
        int s = 0;
        if (*node==NULL) return;
        if ((*node)->left!=NULL) s+=(*node)->left->val;
        if ((*node)->right!=NULL) s+=(*node)->right->val;
        (*node)->val = s==(*node)->val? 0:1;
        reset(&((*node)->left));
        reset(&((*node)->right));
    }
    TreeNode* pruneTree(TreeNode* root) {
        TreeNode** node = &root;
        sum(node);
        chk(node);
        reset(node);
        return *node;
    }
};