Codeforces Round #382 (Div. 2), problem: (E) Ostap and Tree, Accepted Solution in C/C++

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#define M	1000000007

int add(int a, int b) { return (a + b) % M; }
int mult(int a, int b) { return ((long long) a * b) % M; }

#define K	20

struct edge {
	int v;
	struct edge *next;
};

struct vertex {
	struct edge *list;
	int dp[K + 1][K + 2];
} *vv;

int max(int a, int b, int k) {
	return a != k + 1 && b != k + 1 ? (a > b ? a : b) : (a == k + 1 ? b : a);
}

void edgeadd(int u, int v) {
	struct edge *x = malloc(sizeof(*x));

	x->v = v;
	x->next = vv[u].list;
	vv[u].list = x;
}

void dfs(int u, int v, int n, int k) {
	int i, j, wi, wj, dp[K + 1][K + 2];
	struct edge *x;

	for (x = vv[v].list; x != NULL; x = x->next)
		if (x->v != u) 
			dfs(v, x->v, n, k);
	vv[v].dp[0][k + 1] = vv[v].dp[k][0] = 1; /* the jth node */
	for (x = vv[v].list; x != NULL; x = x->next) { /* the tree */
		int w = x->v;

		memset(dp, 0, (K + 1) * (K + 2) * sizeof(**dp));
		if (w != u) {
			for (i = 0; i <= k; i++)
				for (j = 0; j <= k + 1; j++)
					if (vv[v].dp[i][j] != 0)
						for (wi = 0; wi <= k; wi++)
							for (wj = 0; wj <= k + 1; wj++)
								if (vv[w].dp[wi][wj] != 0) {
									int i_ = i < wi + 1 ? i : wi + 1;
									int j_ = k + 1;

									if (j != k + 1 && j + wi + 1 > k)
										j_ = max(j_, j, k);
									if (wj != k + 1 && wj + 1 + i > k)
										j_ = max(j_, wj + 1, k);
									if (j_ == k)
										continue;
									dp[i_][j_] = add(dp[i_][j_], mult(vv[v].dp[i][j], vv[w].dp[wi][wj]));
								}
			memcpy(vv[v].dp, dp, sizeof(**dp) * (K + 1) * (K + 2));
		}
	}
}

int main() {
	int i, n, k, cnt;

	scanf("%d%d", &n, &k);
	vv = calloc(n, sizeof(*vv));
	if (k == 0)
		printf("1\n");
	else {
		for (i = 1; i < n; i++) {
			int u, v;

			scanf("%d%d", &u, &v);
			u--, v--;
			edgeadd(u, v);
			edgeadd(v, u);
		}
		dfs(-1, 0, n, k);
		cnt = 0;
		for (i = 0; i <= k; i++)
			cnt = add(cnt, vv[0].dp[i][k + 1]);
		printf("%d\n", cnt);
	}
	return 0;
}

#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#define M	1000000007

int add(int a, int b) { return (a + b) % M; }
int mult(int a, int b) { return ((long long) a * b) % M; }

#define K	20

struct edge {
	int v;
	struct edge *next;
};

struct vertex {
	struct edge *list;
	int dp[K + 1][K + 2];
} *vv;

int max(int a, int b, int k) {
	return a != k + 1 && b != k + 1 ? (a > b ? a : b) : (a == k + 1 ? b : a);
}

void edgeadd(int u, int v) {
	struct edge *x = malloc(sizeof(*x));

	x->v = v;
	x->next = vv[u].list;
	vv[u].list = x;
}

void dfs(int u, int v, int n, int k) {
	int i, j, wi, wj, dp[K + 1][K + 2];
	struct edge *x;

	for (x = vv[v].list; x != NULL; x = x->next)
		if (x->v != u) 
			dfs(v, x->v, n, k);
	vv[v].dp[0][k + 1] = vv[v].dp[k][0] = 1; /* the jth node */
	for (x = vv[v].list; x != NULL; x = x->next) { /* the tree */
		int w = x->v;

		memset(dp, 0, (K + 1) * (K + 2) * sizeof(**dp));
		if (w != u) {
			for (i = 0; i <= k; i++)
				for (j = 0; j <= k + 1; j++)
					if (vv[v].dp[i][j] != 0)
						for (wi = 0; wi <= k; wi++)
							for (wj = 0; wj <= k + 1; wj++)
								if (vv[w].dp[wi][wj] != 0) {
									int i_ = i < wi + 1 ? i : wi + 1;
									int j_ = k + 1;

									if (j != k + 1 && j + wi + 1 > k)
										j_ = max(j_, j, k);
									if (wj != k + 1 && wj + 1 + i > k)
										j_ = max(j_, wj + 1, k);
									if (j_ == k)
										continue;
									dp[i_][j_] = add(dp[i_][j_], mult(vv[v].dp[i][j], vv[w].dp[wi][wj]));
								}
			memcpy(vv[v].dp, dp, sizeof(**dp) * (K + 1) * (K + 2));
		}
	}
}

int main() {
	int i, n, k, cnt;

	scanf("%d%d", &n, &k);
	vv = calloc(n, sizeof(*vv));
	if (k == 0)
		printf("1\n");
	else {
		for (i = 1; i < n; i++) {
			int u, v;

			scanf("%d%d", &u, &v);
			u--, v--;
			edgeadd(u, v);
			edgeadd(v, u);
		}
		dfs(-1, 0, n, k);
		cnt = 0;
		for (i = 0; i <= k; i++)
			cnt = add(cnt, vv[0].dp[i][k + 1]);
		printf("%d\n", cnt);
	}
	return 0;
}