Advertisement
qwerty787788

heavy-light && treap

Jul 18th, 2015
290
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Java 10.67 KB | None | 0 0
  1. import java.io.*;
  2. import java.util.*;
  3.  
  4. public class TestClass {
  5.     FastScanner in;
  6.     PrintWriter out;
  7.  
  8.     Random rnd = new Random(123);
  9.  
  10.     class Node {
  11.         long priority = rnd.nextLong();
  12.         int value, sum, sz;
  13.         boolean needReverse;
  14.         Node left, right;
  15.  
  16.         Node(int value) {
  17.             this.value = value;
  18.             this.sz = 1;
  19.             this.sum = value;
  20.         }
  21.  
  22.         @Override
  23.         public String toString() {
  24.             return "Node [priority=" + priority + ", value=" + value + ", sum="
  25.                     + sum + ", sz=" + sz + ", needReverse=" + needReverse + "]";
  26.         }
  27.  
  28.         void updateValues() {
  29.             sum = value + sum(left) + sum(right);
  30.             sz = 1 + sz(left) + sz(right);
  31.         }
  32.  
  33.     }
  34.  
  35.     Node merge(Node f, Node s) {
  36.         if (f == null || s == null) {
  37.             return f == null ? s : f;
  38.         }
  39.         upd(f);
  40.         upd(s);
  41.         if (f.priority < s.priority) {
  42.             f.right = merge(f.right, s);
  43.             f.updateValues();
  44.             return f;
  45.         } else {
  46.             s.left = merge(f, s.left);
  47.             s.updateValues();
  48.             return s;
  49.         }
  50.     }
  51.  
  52.     Node[] split(Node node, int leftSize) {
  53.         if (leftSize == 0) {
  54.             return new Node[] { null, node };
  55.         }
  56.         upd(node);
  57.         if (sz(node.left) >= leftSize) {
  58.             Node[] tmp = split(node.left, leftSize);
  59.             node.left = tmp[1];
  60.             node.updateValues();
  61.             return new Node[] { tmp[0], node };
  62.         }
  63.         Node[] tmp = split(node.right, leftSize - sz(node.left) - 1);
  64.         node.right = tmp[0];
  65.         node.updateValues();
  66.         return new Node[] { node, tmp[1] };
  67.     }
  68.  
  69.     int sz(Node n) {
  70.         if (n == null) {
  71.             return 0;
  72.         }
  73.         return n.sz;
  74.     }
  75.  
  76.     int sum(Node n) {
  77.         if (n == null) {
  78.             return 0;
  79.         }
  80.         return n.sum;
  81.     }
  82.  
  83.     void upd(Node n) {
  84.         if (n.needReverse) {
  85.             Node tmp = n.left;
  86.             n.left = n.right;
  87.             n.right = tmp;
  88.             n.needReverse = false;
  89.             if (n.left != null) {
  90.                 n.left.needReverse = !n.left.needReverse;
  91.             }
  92.             if (n.right != null) {
  93.                 n.right.needReverse = !n.right.needReverse;
  94.             }
  95.         }
  96.     }
  97.  
  98.     Node reverse(Node n, int from, int to) {
  99.         Node[] tmp = split(n, from);
  100.         Node[] tmp2 = split(tmp[1], to - from + 1);
  101.         tmp2[0].needReverse = true;
  102.         return merge(tmp[0], merge(tmp2[0], tmp2[1]));
  103.     }
  104.  
  105.     int getSum(Node n, int size) {
  106.         if (n == null) {
  107.             return 0;
  108.         }
  109.         upd(n);
  110.         if (size >= sz(n)) {
  111.             return sum(n);
  112.         }
  113.         if (size <= 0) {
  114.             return 0;
  115.         }
  116.         int res = getSum(n.left, size);
  117.         size -= sz(n.left);
  118.         if (size > 0) {
  119.             res += n.value;
  120.         }
  121.         res += getSum(n.right, size - 1);
  122.         return res;
  123.     }
  124.  
  125.     ArrayList<Integer>[] g;
  126.     int[] h;
  127.     int[] sz;
  128.     int[][] up;
  129.     final int MAX = 20;
  130.  
  131.     int[] q;
  132.  
  133.     void go(int vv, int par) {
  134.         int n = sz.length;
  135.         q = new int[n];
  136.         q[0] = vv;
  137.         int qIt = 0, qSz = 1;
  138.         up[0][vv] = vv;
  139.         h[vv] = 0;
  140.         while (qIt < qSz) {
  141.             int v = q[qIt++];
  142.             int p = up[0][v];
  143.             for (int i = 1; i < MAX; i++) {
  144.                 up[i][v] = up[i - 1][up[i - 1][v]];
  145.             }
  146.             for (int to : g[v]) {
  147.                 if (to == p) {
  148.                     continue;
  149.                 }
  150.                 h[to] = h[v] + 1;
  151.                 q[qSz++] = to;
  152.                 up[0][to] = v;
  153.             }
  154.         }
  155.         for (int i = n - 1; i >= 0; i--) {
  156.             int v = q[i];
  157.             sz[v] = 1;
  158.             for (int to : g[v]) {
  159.                 if (to != up[0][v]) {
  160.                     sz[v] += sz[to];
  161.                 }
  162.             }
  163.         }
  164.     }
  165.  
  166.     int[] a;
  167.  
  168.     int lca(int x, int y) {
  169.         if (h[x] < h[y]) {
  170.             int tmp = x;
  171.             x = y;
  172.             y = tmp;
  173.         }
  174.         for (int i = MAX - 1; i >= 0; i--) {
  175.             if (h[up[i][x]] >= h[y]) {
  176.                 x = up[i][x];
  177.             }
  178.         }
  179.         for (int i = MAX - 1; i >= 0; i--) {
  180.             if (up[i][x] != up[i][y]) {
  181.                 x = up[i][x];
  182.                 y = up[i][y];
  183.             }
  184.         }
  185.         if (x == y) {
  186.             return x;
  187.         }
  188.         return up[0][x];
  189.     }
  190.  
  191.     int goUp(int v, int h) {
  192.         for (int i = MAX - 1; i >= 0; i--) {
  193.             if (((1 << i) & h) != 0) {
  194.                 v = up[i][v];
  195.             }
  196.         }
  197.         return v;
  198.     }
  199.  
  200.     int[] endOfPath;
  201.     int[] startOfPath;
  202.     Node[] pathNode;
  203.  
  204.     void _go2(int vv, int pp) {
  205.         startOfPath[vv] = vv;
  206.         endOfPath[vv] = vv;
  207.         for (int v : q) {
  208.             startOfPath[v] = v;
  209.             int p = up[0][v];
  210.             for (int to : g[v]) {
  211.                 if (to == p) {
  212.                     continue;
  213.                 }
  214.                 if (sz[to] * 2 >= sz[v]) {
  215.                     endOfPath[to] = endOfPath[v];
  216.                 } else {
  217.                     endOfPath[to] = to;
  218.                 }
  219.             }
  220.         }
  221.         for (int i = q.length - 1; i >= 0; i--) {
  222.             int v = q[i];
  223.             int p = up[0][v];
  224.             for (int to : g[v]) {
  225.                 if (to == p) {
  226.                     continue;
  227.                 }
  228.                 if (sz[to] * 2 >= sz[v]) {
  229.                     startOfPath[v] = startOfPath[to];
  230.                 }
  231.             }
  232.             int root = endOfPath[v];
  233.             pathNode[root] = merge(pathNode[root], new Node(a[v]));
  234.         }
  235.     }
  236.  
  237.     int getSumPathUp(int v, int cnt) {
  238.         if (cnt == 0) {
  239.             return 0;
  240.         }
  241.         int maxH = h[startOfPath[v]], minH = h[endOfPath[v]];
  242.         int hereSize = Math.min(cnt, h[v] - minH + 1);
  243.         int result = 0;
  244.         if (hereSize < cnt) {
  245.             result += getSumPathUp(up[0][endOfPath[v]], cnt - hereSize);
  246.         }
  247.         Node node = pathNode[endOfPath[v]];
  248.         result += getSum(node, maxH - h[v] + hereSize)
  249.                 - getSum(node, maxH - h[v]);
  250.         return result;
  251.     }
  252.  
  253.     int getSum(int v, int u) {
  254.         int lca = lca(u, v);
  255.         return getSumPathUp(v, h[v] - h[lca] + 1)
  256.                 + getSumPathUp(u, h[u] - h[lca]);
  257.     }
  258.  
  259.     void generatePath(int v, int cnt, ArrayList<Integer> allV,
  260.             ArrayList<Node> leftPart, ArrayList<Node> rightPart,
  261.             ArrayList<Node> mainPart, ArrayList<Integer> sizes) {
  262.         if (cnt == 0) {
  263.             return;
  264.         }
  265.         int maxH = h[startOfPath[v]], minH = h[endOfPath[v]];
  266.         int hereSize = Math.min(cnt, h[v] - minH + 1);
  267.         Node node = pathNode[endOfPath[v]];
  268.         sizes.add(hereSize);
  269.         allV.add(endOfPath[v]);
  270.         Node[] tmp = split(node, maxH - h[v]);
  271.         Node[] tmp2 = split(tmp[1], hereSize);
  272.         leftPart.add(tmp[0]);
  273.         rightPart.add(tmp2[1]);
  274.         mainPart.add(tmp2[0]);
  275.         generatePath(up[0][endOfPath[v]], cnt - hereSize, allV, leftPart,
  276.                 rightPart, mainPart, sizes);
  277.     }
  278.  
  279.     void reverse(int v, int u) {
  280.         int lca = lca(v, u);
  281.         if (u != lca) {
  282.             int tmp = goUp(u, h[u] - h[lca] - 1);
  283.             if (endOfPath[tmp] != tmp) {
  284.                 tmp = v;
  285.                 v = u;
  286.                 u = tmp;
  287.             }
  288.         }
  289.         ArrayList<Integer> allVFirst = new ArrayList<>();
  290.         ArrayList<Integer> sizesFirst = new ArrayList<>();
  291.         ArrayList<Node> leftPartFirst = new ArrayList<>();
  292.         ArrayList<Node> rightPartFirst = new ArrayList<>();
  293.         ArrayList<Node> mainPartFirst = new ArrayList<>();
  294.         generatePath(v, h[v] - h[lca] + 1, allVFirst, leftPartFirst,
  295.                 rightPartFirst, mainPartFirst, sizesFirst);
  296.         ArrayList<Integer> allVSecond = new ArrayList<>();
  297.         ArrayList<Integer> sizesSecond = new ArrayList<>();
  298.         ArrayList<Node> leftPartSecond = new ArrayList<>();
  299.         ArrayList<Node> rightPartSecond = new ArrayList<>();
  300.         ArrayList<Node> mainPartSecond = new ArrayList<>();
  301.         generatePath(u, h[u] - h[lca], allVSecond, leftPartSecond,
  302.                 rightPartSecond, mainPartSecond, sizesSecond);
  303.         Node node = null;
  304.         for (Node x : mainPartFirst) {
  305.             node = merge(node, x);
  306.         }
  307.         for (int i = mainPartSecond.size() - 1; i >= 0; i--) {
  308.             Node x = mainPartSecond.get(i);
  309.             x.needReverse = !x.needReverse;
  310.             node = merge(node, x);
  311.         }
  312.         node.needReverse = !node.needReverse;
  313.         for (int i = 0; i < mainPartFirst.size(); i++) {
  314.             Node[] tmp = split(node, sizesFirst.get(i));
  315.             pathNode[allVFirst.get(i)] = merge(leftPartFirst.get(i),
  316.                     merge(tmp[0], rightPartFirst.get(i)));
  317.             node = tmp[1];
  318.         }
  319.         for (int i = mainPartSecond.size() - 1; i >= 0; i--) {
  320.             Node[] tmp = split(node, sizesSecond.get(i));
  321.             tmp[0].needReverse = !tmp[0].needReverse;
  322.             pathNode[allVSecond.get(i)] = merge(leftPartSecond.get(i),
  323.                     merge(tmp[0], rightPartSecond.get(i)));
  324.             node = tmp[1];
  325.         }
  326.     }
  327.  
  328.     void solve() {
  329.         int n = in.nextInt();
  330.         int qq = in.nextInt();
  331.         a = new int[n];
  332.         g = new ArrayList[n];
  333.         for (int i = 0; i < n; i++) {
  334.             a[i] = in.nextInt();
  335.             g[i] = new ArrayList<>();
  336.         }
  337.         for (int i = 0; i + 1 < n; i++) {
  338.             int fr = in.nextInt() - 1, to = in.nextInt() - 1;
  339.             g[fr].add(to);
  340.             g[to].add(fr);
  341.         }
  342.         h = new int[n];
  343.         h[0] = -1;
  344.         sz = new int[n];
  345.         up = new int[MAX][n];
  346.         go(0, 0);
  347.         pathNode = new Node[n];
  348.         startOfPath = new int[n];
  349.         endOfPath = new int[n];
  350.         _go2(0, 0);
  351.         for (int q = 0; q < qq; q++) {
  352.             if (in.next().equals("R")) {
  353.                 int u = in.nextInt() - 1, v = in.nextInt() - 1;
  354.                 reverse(u, v);
  355.             } else {
  356.                 int u = in.nextInt() - 1, v = in.nextInt() - 1;
  357.                 out.println(getSum(u, v));
  358.             }
  359.         }
  360.     }
  361.  
  362.     void solve3333() {
  363.         for (int iter = 0; iter < 123123; iter++) {
  364.             System.err.println("iter = " + iter);
  365.             if (iter == 28) {
  366.                 System.err.println("???");
  367.             }
  368.             solve123();
  369.         }
  370.     }
  371.  
  372.     void solve123() {
  373.         final int MMAX = 100000;
  374.         int n = MMAX;
  375.         int qq = MMAX;
  376.         a = new int[n];
  377.         g = new ArrayList[n];
  378.         for (int i = 0; i < n; i++) {
  379.             a[i] = rnd.nextInt(1000);
  380.             g[i] = new ArrayList<>();
  381.         }
  382.         for (int i = 0; i + 1 < n; i++) {
  383.             // int fr = in.nextInt() - 1, to = in.nextInt() - 1;
  384.             int fr = i + 1, to = rnd.nextInt(fr);
  385.             // int fr = i, to = i + 1;
  386.             g[fr].add(to);
  387.             g[to].add(fr);
  388.         }
  389.         h = new int[n];
  390.         h[0] = -1;
  391.         sz = new int[n];
  392.         up = new int[MAX][n];
  393.         go(0, 0);
  394.         pathNode = new Node[n];
  395.         startOfPath = new int[n];
  396.         endOfPath = new int[n];
  397.         _go2(0, 0);
  398.         for (int q = 0; q < qq; q++) {
  399.             if (rnd.nextBoolean()) {
  400.                 int u = rnd.nextInt(n), v = rnd.nextInt(n);
  401.                 reverse(u, v);
  402.             } else {
  403.                 int u = rnd.nextInt(n), v = rnd.nextInt(n);
  404.                 // out.println(getSum(u, v));
  405.                 getSum(u, v);
  406.             }
  407.         }
  408.     }
  409.  
  410.     void run() {
  411.         try {
  412.             in = new FastScanner(new File("object.in"));
  413.             out = new PrintWriter(new File("object.out"));
  414.  
  415.             solve();
  416.  
  417.             out.close();
  418.         } catch (FileNotFoundException e) {
  419.             e.printStackTrace();
  420.         }
  421.     }
  422.  
  423.     void runIO() {
  424.  
  425.         in = new FastScanner(System.in);
  426.         out = new PrintWriter(System.out);
  427.  
  428.         solve();
  429.  
  430.         out.close();
  431.     }
  432.  
  433.     class FastScanner {
  434.         BufferedReader br;
  435.         StringTokenizer st;
  436.  
  437.         public FastScanner(File f) {
  438.             try {
  439.                 br = new BufferedReader(new FileReader(f));
  440.             } catch (FileNotFoundException e) {
  441.                 e.printStackTrace();
  442.             }
  443.         }
  444.  
  445.         public FastScanner(InputStream f) {
  446.             br = new BufferedReader(new InputStreamReader(f));
  447.         }
  448.  
  449.         String next() {
  450.             while (st == null || !st.hasMoreTokens()) {
  451.                 String s = null;
  452.                 try {
  453.                     s = br.readLine();
  454.                 } catch (IOException e) {
  455.                     e.printStackTrace();
  456.                 }
  457.                 if (s == null)
  458.                     return null;
  459.                 st = new StringTokenizer(s);
  460.             }
  461.             return st.nextToken();
  462.         }
  463.  
  464.         boolean hasMoreTokens() {
  465.             while (st == null || !st.hasMoreTokens()) {
  466.                 String s = null;
  467.                 try {
  468.                     s = br.readLine();
  469.                 } catch (IOException e) {
  470.                     e.printStackTrace();
  471.                 }
  472.                 if (s == null)
  473.                     return false;
  474.                 st = new StringTokenizer(s);
  475.             }
  476.             return true;
  477.         }
  478.  
  479.         int nextInt() {
  480.             return Integer.parseInt(next());
  481.         }
  482.  
  483.         long nextLong() {
  484.             return Long.parseLong(next());
  485.         }
  486.  
  487.         double nextDouble() {
  488.             return Double.parseDouble(next());
  489.         }
  490.     }
  491.  
  492.     public static void main(String[] args) {
  493.         new TestClass().runIO();
  494.     }
  495. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement