1261. Find Elements in a Contaminated Binary Tree

, , ,

DFS + Bit

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
class FindElements {

    private final long[] bitmap;

    public FindElements(TreeNode root) {
        bitmap = new long[1048576 / 64];
        root.val = 0;
        // 不要忘记标记 root 节点
        markBit(bitIndexToArrayIndex(root.val), bitIndexToOffset(root.val));
        dfs(root.left, root, true);
        dfs(root.right, root, false);
    }

    private void dfs(TreeNode node, TreeNode parent, boolean isLeft) {
        if (node == null) {
            return;
        }

        if (isLeft) {
            node.val = parent.val * 2 + 1;
        } else {
            node.val = parent.val * 2 + 2;
        }
        markBit(bitIndexToArrayIndex(node.val), bitIndexToOffset(node.val));
        dfs(node.left, node, true);
        dfs(node.right, node, false);
    }

    private void markBit(int arrayIndex, int offset) {
        bitmap[arrayIndex] |= 1L << offset;
    }

    private int bitIndexToArrayIndex(int bitIndex) {
        return bitIndex / 64;
    }

    private int bitIndexToOffset(int bitIndex) {
        return bitIndex % 64;
    }

    public boolean find(int target) {
        int arrayIndex = bitIndexToArrayIndex(target);
        int offset = bitIndexToOffset(target);

        return (bitmap[arrayIndex] & 1L << offset) != 0;
    }

}
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
class FindElements {

    private final TreeNode root;

    public FindElements(TreeNode root) {
        this.root = root;
        root.val = 0;
        dfs(root.left, root, true);
        dfs(root.right, root, false);
    }

    private void dfs(TreeNode node, TreeNode parent, boolean isLeft) {
        if (node == null) {
            return;
        }

        if (isLeft) {
            node.val = parent.val * 2 + 1;
        } else {
            node.val = parent.val * 2 + 2;
        }
        dfs(node.left, node, true);
        dfs(node.right, node, false);
    }

    public boolean find(int target) {
        target++; // 想象为树中所有节点值加 1
        TreeNode node = root;
        int mask = Integer.highestOneBit(target) >> 1; // 注意从次高位开始，因为 root 已经为 1
        while (mask != 0 && node != null) {
            if ((target & mask) == 0) {
                node = node.left;
            } else {
                node = node.right;
            }

            mask >>= 1;
        }

        return node != null;
    }

}

References

1261. Find Elements in a Contaminated Binary Tree