Leetcode/src/main/java/com/fishercoder/solutions/secondthousand/_1530.java at master · fishercoder1534/Leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
package com.fishercoder.solutions.secondthousand;

import com.fishercoder.common.classes.TreeNode;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

public class _1530 {
    public static class Solution1 {
        public int countPairs(TreeNode root, int distance) {
            Map<TreeNode, TreeNode> childToParentMap = new HashMap<>();
            List<TreeNode> leafNodes = new ArrayList();
            postOrderToFindAllLeavesAndBuildChildParentMap(root, childToParentMap, leafNodes);
            int pairs = 0;
            for (TreeNode leaf : leafNodes) {
                pairs += bfsToPossibleLeaves(leaf, distance, childToParentMap);
            }
            return pairs / 2;
        }

        private int bfsToPossibleLeaves(
                TreeNode leaf, int distance, Map<TreeNode, TreeNode> childToParentMap) {
            Queue<TreeNode> q = new LinkedList<>();
            q.offer(leaf);
            Set<TreeNode> visited = new HashSet<>();
            visited.add(leaf);
            int count = 0;
            while (!q.isEmpty() && distance >= 0) {
                int size = q.size();
                for (int i = 0; i < size; i++) {
                    TreeNode curr = q.poll();
                    if (leaf != curr && curr.left == null && curr.right == null) {
                        count++;
                    }
                    if (curr.left != null && visited.add(curr.left)) {
                        q.offer(curr.left);
                    }
                    if (curr.right != null && visited.add(curr.right)) {
                        q.offer(curr.right);
                    }
                    if (childToParentMap.containsKey(curr)
                            && visited.add(childToParentMap.get(curr))) {
                        q.offer(childToParentMap.get(curr));
                    }
                }
                distance--;
            }
            return count;
        }

        private void postOrderToFindAllLeavesAndBuildChildParentMap(
                TreeNode node, Map<TreeNode, TreeNode> childToParentMap, List<TreeNode> leafNodes) {
            if (node == null) {
                return;
            }
            if (node.left != null) {
                childToParentMap.put(node.left, node);
            }
            if (node.right != null) {
                childToParentMap.put(node.right, node);
            }
            postOrderToFindAllLeavesAndBuildChildParentMap(node.left, childToParentMap, leafNodes);
            postOrderToFindAllLeavesAndBuildChildParentMap(node.right, childToParentMap, leafNodes);
            if (node.left == null && node.right == null) {
                leafNodes.add(node);
            }
        }
    }
}