算法实现:一次dfs + 求和indexed tree的多次update,query完成

    1. 由于要查找所有祖先节点的值,比自己大的个数,dfs过程中,祖先节点可大,可小求和有困难,所以使用indexed tree来求和
    2. 对每个节点来说其实是求一个自身到根节点的区间中比自己大的个数的和
    3. 并且为了只走一次dfs,所以当进入V节点时,update(V,1),离开节点时update(V,-1),保证了祖先节点不会重复计算
    4. 注意:栈内存给3M的话,20个case都能通过,本社那边的case可能数据量较大,所以把dfs修改为使用stack的非递归实现

    测试用例:
    2
    7 7
    7 4
    6 2
    7 5
    6 1
    1 3
    7 6
    10 2
    1 3
    5 4
    5 8
    5 9
    9 1
    2 5
    9 6
    1 10
    5 7
    (输出)
    #1 9
    #2 7

    1. public class 路径管理 {
    2. private static final HashMap<Integer, LinkedList<Integer>> adj = new HashMap<>();
    3. private static final LinkedList<Integer> stack = new LinkedList<>();
    4. private static final long[] tree = new long[260_005];
    5. private static int N, start, offset;
    6. public static void main(String[] args) throws Exception {
    7. StreamTokenizer in = new StreamTokenizer(new BufferedReader(new InputStreamReader(System.in)));
    8. in.nextToken();
    9. int T = (int) in.nval;
    10. for (int t = 1; t <= T; t++) {
    11. in.nextToken();
    12. N = (int) in.nval;
    13. in.nextToken();
    14. start = (int) in.nval;
    15. Arrays.fill(tree, 0);
    16. for (int i = 1; i < N; i++) {
    17. in.nextToken();
    18. int a = (int) in.nval;
    19. in.nextToken();
    20. int b = (int) in.nval;
    21. adj.computeIfAbsent(a, l -> new LinkedList<>()).add(b);
    22. }
    23. offset = 1;
    24. while (offset < N)
    25. offset <<= 1;
    26. offset--;
    27. System.out.println("#" + t + " " + dfs());
    28. }
    29. }
    30. private static long dfs() {
    31. stack.add(start);
    32. update(start, 1);
    33. long sum = 0;
    34. while (!stack.isEmpty()) {
    35. LinkedList<Integer> next = adj.get(stack.peekLast());
    36. if (next == null || next.isEmpty()) {
    37. int out = stack.pollLast();
    38. update(out, -1);
    39. } else {
    40. int in = next.poll();
    41. sum += query(in + 1, N);
    42. stack.add(in);
    43. update(in, 1);
    44. }
    45. }
    46. return sum;
    47. }
    48. private static void update(int i, int v) {
    49. i = offset + i;
    50. while (i > 0) {
    51. tree[i] += v;
    52. i >>= 1;
    53. }
    54. }
    55. private static long query(int s, int e) {
    56. long sum = 0;
    57. s += offset;
    58. e += offset;
    59. while (s <= e) {
    60. if (s % 2 == 1)
    61. sum += tree[s];
    62. if (e % 2 == 0)
    63. sum += tree[e];
    64. s = (s + 1) >> 1;
    65. e = (e - 1) >> 1;
    66. }
    67. return sum;
    68. }
    69. }