git reimport

archive/2017.03/2017.03.18 - VK Cup 2017 - Round 1/TaskC.java

@@ -0,0 +1,128 @@
+package chelper;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+import io.InputReader;
+import io.OutputWriter;
+
+public class TaskC {
+  int n, k;
+  List<List<Integer>> edges;
+
+  long ans = 0;
+
+  long[] rotCount(long[] x) {
+    long[] res = new long[k];
+    for (int i = 0; i < k; i++) {
+      res[i] = x[(i - 1 + k) % k];
+    }
+    return res;
+  }
+
+  long[] rotAcc(long[] x) {
+    long[] res = new long[k];
+    for (int i = 0; i < k; i++) {
+      res[i] = x[(i - 1 + k) % k];
+    }
+    return res;
+  }
+
+  long[] copy(long[] x) {
+    return Arrays.copyOf(x, k);
+  }
+
+  void add(long[] a, long[] b) {
+    for (int i = 0; i < k; i++) {
+      a[i] += b[i];
+    }
+  }
+
+  public void solve(int testNumber, InputReader in, OutputWriter out) {
+    n = in.nextInt();
+    k = in.nextInt();
+
+    visited = new boolean[n];
+
+    edges = new ArrayList<>();
+    for (int i = 0; i < n; i++) {
+      edges.add(new ArrayList<>());
+    }
+
+    for (int i = 0; i < n - 1; i++) {
+      int a = in.nextInt() - 1;
+      int b = in.nextInt() - 1;
+
+      edges.get(a).add(b);
+      edges.get(b).add(a);
+    }
+
+//    long[] count = new long[k];
+    long[] acc = new long[k];
+
+    Result ans = dfs(0, 0, acc);
+
+    out.println(ans.cost / 2);
+  }
+
+  class Result {
+    long cost;
+    long[] acc;
+
+    public Result(long cost, long[] acc) {
+      this.cost = cost;
+      this.acc = acc;
+    }
+  }
+
+  boolean[] visited;
+
+//  void addComp(long[] a, long[] b) {
+//    for (int i = 0; i < k; i++) {
+//      for (int j = 0; j < k; j++) {
+//        if (i == j) {
+//          continue;
+//        }
+//        long t = a[i] + b[j];
+//        if (i + j >= k) {
+//          ans += ;
+//        }
+//      }
+//    }
+//  }
+
+  Result dfs(int v, long parentCost, long[] accParent) {
+    visited[v] = true;
+
+    long costParentChildren = parentCost;
+    long costChildren = 0;
+
+    long[] accParentChildren = new long[k];
+    costParentChildren += rot(accParentChildren, accParent);
+    accParentChildren[0] += 1;
+
+    long[] accChildren = new long[k];
+    accChildren[0] += 1;
+
+    for (Integer child : edges.get(v)) {
+      if (visited[child]) {
+        continue;
+      }
+
+      Result atChild = dfs(child, costParentChildren, accParentChildren);
+
+      long[] childAcc = new long[k];
+
+      costChildren += atChild.cost;
+      costChildren += rot(childAcc, atChild.acc);
+      costParentChildren += atChild.cost;
+      costParentChildren += rot(childAcc, atChild.acc);
+
+      add(accParentChildren, childAcc);
+      add(accChildren, childAcc);
+    }
+
+    return new Result(costChildren, accChildren);
+  }
+}