推特

Posted on In Disqus:

迎难而上，祝我好运。面朝大海，春暖花开。

OA

  • 总共有n个节点，给两个List表示从first到second有一条undirected edge，再给一个values数组表示对应编号的node的值。然后给一堆query，找这些编号为root的子树中素数的节点count。注意编号为1的节点作为根。
  • 直接把自己做的存下来了：利用1作为根这个条件用BFS构建多叉树，然后用递归count有多少个素数。为了避免重复计算，用了memo（一个用在判断prime，一个用在计数prime node）.
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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
static class TreeNode {
    int id;
    int val;
    List<TreeNode> children;
    public TreeNode(int id, int val) {
        this.id = id;
        this.val = val;
        children = new ArrayList<>();
    }
}
// Complete the primeQuery function below.
static List<Integer> primeQuery(int n, List<Integer> first, List<Integer> second, List<Integer> values, List<Integer> queries) {
    if (first == null || second == null || first.size() != second.size() || values.size() != n) {
        return null;
    }
    Map<Integer, Set<Integer>> edgeMap = new HashMap<>();   // map from node-id to a set of its neighbors' id
    for (int i = 0; i < first.size(); i++) {
        edgeMap.putIfAbsent(first.get(i), new HashSet<>());
        edgeMap.putIfAbsent(second.get(i), new HashSet<>());
        edgeMap.get(first.get(i)).add(second.get(i));       // undirected: add the mapping relation to both sides
        edgeMap.get(second.get(i)).add(first.get(i));
    }
    
    Map<Integer, TreeNode> id2Node = new HashMap<>();       // map from id to actual TreeNode object
    TreeNode root = constructTree(edgeMap, values, id2Node);
    List<Integer> ans = new ArrayList<>();
    Map<Integer, Integer> countCache = new HashMap<>();
    Map<Integer, Boolean> primeCache = new HashMap<>();
    for (int query : queries) {
        ans.add(getPrimeCount(id2Node.get(query), countCache, primeCache));
    }
    return ans;
}

// use BFS to constrct tree
static TreeNode constructTree(Map<Integer, Set<Integer>> edgeMap, List<Integer> values, 
                                Map<Integer, TreeNode> id2Node) {
    Queue<TreeNode> q = new LinkedList<>();
    TreeNode root = new TreeNode(1, values.get(0));     // start from root with id 1
    id2Node.put(1, root);
    q.add(root);
    while (!q.isEmpty()) {
        TreeNode curr = q.poll();
        Set<Integer> childIds = edgeMap.get(curr.id);     // get all its neighbors' id
        if (childIds == null) {
            continue;
        }
        for (int childId : childIds) {
            TreeNode child = new TreeNode(childId, values.get(childId - 1));
            id2Node.put(childId, child);                        // from id to node
            q.offer(child);
            curr.children.add(child);                            // add neighbor node into curr node's children list
            edgeMap.get(childId).remove(curr.id);       // remove curr's id from the neighbor's "neighbor set"
        }
        edgeMap.remove(curr.id);                            // finish process curr node, remove it from edgeMap
    }
    return root;
}

// given a TreeNode, return the count of nodes with prime value
static int getPrimeCount(TreeNode node, Map<Integer, Integer> countCache, 
                            Map<Integer, Boolean> primeCache) {
    if (node == null) {
        return 0;
    }
    if (countCache.containsKey(node.id)) {  // hit in cache, simply return
        return countCache.get(node.id);
    }
    int count = isPrime(node.val, primeCache) ? 1 : 0;
    for (TreeNode child : node.children) {   // get prime count at each child
        count += getPrimeCount(child, countCache, primeCache);
    }
    countCache.put(node.id, count);     // put into cache
    return count;
}

static boolean isPrime(int val, Map<Integer, Boolean> primeCache) {
    if (primeCache.containsKey(val)) {   // hit in cache, simply return
        return primeCache.get(val);
    }
    if (val <= 1) {
        primeCache.put(val, false);
        return false;
    }
    int limit = (int) Math.sqrt(val);
    for (int i = 2; i <= limit; i++) {
        if (val % i == 0) {
            primeCache.put(val, false);
            return false;
        }
    }
    primeCache.put(val, true);
    return true;
}