modify code

Author: algorithmzuo

src/class183/Code05_Message1.java

@@ -30,7 +30,6 @@ public class Code05_Message1 {
 	public static int[] maxPart = new int[MAXN];
 	public static int centroid;
 
-	public static int[] dep = new int[MAXN];
 	public static int[] depCnt = new int[MAXN];
 	public static int maxDeep;
 
@@ -41,21 +40,23 @@ public class Code05_Message1 {
 	public static int[] ans = new int[MAXN];
 
 	// 讲解118，递归函数改成迭代所需要的栈
-	public static int[][] stack = new int[MAXN][3];
-	public static int stacksize, u, f, e;
+	public static int[][] stack = new int[MAXN][4];
+	public static int stacksize, u, f, dep, e;
 
-	public static void push(int u, int f, int e) {
+	public static void push(int u, int f, int dep, int e) {
 		stack[stacksize][0] = u;
 		stack[stacksize][1] = f;
-		stack[stacksize][2] = e;
+		stack[stacksize][2] = dep;
+		stack[stacksize][3] = e;
 		stacksize++;
 	}
 
 	public static void pop() {
 		--stacksize;
 		u = stack[stacksize][0];
 		f = stack[stacksize][1];
-		e = stack[stacksize][2];
+		dep = stack[stacksize][2];
+		e = stack[stacksize][3];
 	}
 
 	public static void addEdge(int u, int v) {
@@ -92,7 +93,7 @@ public static void getCentroid1(int u, int fa) {
 	// 找重心的迭代版
 	public static void getCentroid2(int cur, int fa) {
 		stacksize = 0;
-		push(cur, fa, -1);
+		push(cur, fa, 0, -1);
 		while (stacksize > 0) {
 			pop();
 			if (e == -1) {
@@ -103,10 +104,10 @@ public static void getCentroid2(int cur, int fa) {
 				e = nextg[e];
 			}
 			if (e != 0) {
-				push(u, f, e);
+				push(u, f, 0, e);
 				int v = tog[e];
 				if (v != f && !vis[v]) {
-					push(tog[e], u, -1);
+					push(tog[e], u, 0, -1);
 				}
 			} else {
 				for (int ei = headg[u]; ei > 0; ei = nextg[ei]) {
@@ -125,37 +126,35 @@ public static void getCentroid2(int cur, int fa) {
 	}
 
 	// 收集信息递归版，java会爆栈，C++可以通过
-	public static void dfs1(int u, int fa) {
-		dep[u] = dep[fa] + 1;
-		depCnt[dep[u]]++;
-		maxDeep = Math.max(maxDeep, dep[u]);
+	public static void dfs1(int u, int fa, int deep) {
+		depCnt[deep]++;
+		maxDeep = Math.max(maxDeep, deep);
 		for (int e = headq[u]; e > 0; e = nextq[e]) {
-			if (tim[e] + 1 >= dep[u]) {
-				timArr[++cnta] = tim[e] - dep[u] + 2;
+			if (tim[e] + 1 >= deep) {
+				timArr[++cnta] = tim[e] - deep + 2;
 				qidArr[cnta] = qid[e];
 			}
 		}
 		for (int e = headg[u]; e > 0; e = nextg[e]) {
 			int v = tog[e];
 			if (v != fa && !vis[v]) {
-				dfs1(v, u);
+				dfs1(v, u, deep + 1);
 			}
 		}
 	}
 
 	// 收集信息迭代版
-	public static void dfs2(int cur, int fa) {
+	public static void dfs2(int cur, int fa, int deep) {
 		stacksize = 0;
-		push(cur, fa, -1);
+		push(cur, fa, deep, -1);
 		while (stacksize > 0) {
 			pop();
 			if (e == -1) {
-				dep[u] = dep[f] + 1;
-				depCnt[dep[u]]++;
-				maxDeep = Math.max(maxDeep, dep[u]);
+				depCnt[dep]++;
+				maxDeep = Math.max(maxDeep, dep);
 				for (int e = headq[u]; e > 0; e = nextq[e]) {
-					if (tim[e] + 1 >= dep[u]) {
-						timArr[++cnta] = tim[e] - dep[u] + 2;
+					if (tim[e] + 1 >= dep) {
+						timArr[++cnta] = tim[e] - dep + 2;
 						qidArr[cnta] = qid[e];
 					}
 				}
@@ -164,10 +163,10 @@ public static void dfs2(int cur, int fa) {
 				e = nextg[e];
 			}
 			if (e != 0) {
-				push(u, f, e);
+				push(u, f, dep, e);
 				int v = tog[e];
 				if (v != f && !vis[v]) {
-					push(tog[e], u, -1);
+					push(tog[e], u, dep + 1, -1);
 				}
 			}
 		}
@@ -176,8 +175,8 @@ public static void dfs2(int cur, int fa) {
 	public static void calc(int u) {
 		cnta = 0;
 		maxDeep = 0;
-		// dfs1(u, 0);
-		dfs2(u, 0);
+		// dfs1(u, 0, 1);
+		dfs2(u, 0, 1);
 		for (int i = 1; i <= cnta; i++) {
 			ans[qidArr[i]] += depCnt[timArr[i]];
 		}
@@ -189,8 +188,8 @@ public static void calc(int u) {
 			if (!vis[v]) {
 				cnta = 0;
 				maxDeep = 0;
-				// dfs1(v, u);
-				dfs2(v, u);
+				// dfs1(v, u, 2);
+				dfs2(v, u, 2);
 				for (int i = 1; i <= cnta; i++) {
 					ans[qidArr[i]] -= depCnt[timArr[i]];
 				}