推特

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

OA

• 总共有n个节点，给两个List表示从first到second有一条undirected edge，再给一个values数组表示对应编号的node的值。然后给一堆query，找这些编号为root的子树中素数的节点count。注意编号为1的节点作为根。
• 直接把自己做的存下来了：利用1作为根这个条件用BFS构建多叉树，然后用递归count有多少个素数。为了避免重复计算，用了memo（一个用在判断prime，一个用在计数prime node）.

  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;
  }