129 lines
2.5 KiB
Java
129 lines
2.5 KiB
Java
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);
|
|
}
|
|
}
|