#include <iostream> #include <vector> #define MOD 1000000007 struct Tree { std::vector<int> p, l, r, s; int read(int n) { p.resize(n); l.resize(n, -1); r.resize(n, -1); s.resize(n); int root = -1; for (int i = 0; i < n; ++i) { std::cin >> p[i]; --p[i]; if (p[i] < 0) { p[i] = -1; root = i; }else { if (i < p[i]) l[p[i]] = i; else r[p[i]] = i; } } calc_size(root); return root; } int calc_size(int i) { int size = 1; if (r[i] != -1) size += calc_size(r[i]); if (l[i] != -1) size += calc_size(l[i]); s[i] = size; return size; } }; int n; Tree A, B; long long stright_left(Tree const& T, int node); long long stright_right(Tree const& T, int node) { long long ret = 0; if (T.r[node] != -1) { ret += stright_right(T, T.r[node]); } if (T.l[node] != -1) { ret += stright_left(T, T.l[node]) + T.s[T.l[node]]; } return ret % MOD; } long long stright_left(Tree const& T, int node) { long long ret = 0; if (T.r[node] != -1) { ret += stright_right(T, T.r[node]) + T.s[T.r[node]]; } if (T.l[node] != -1) { ret += stright_left(T, T.l[node]); } return ret % MOD; } long long reroot(int A_root, int B_root, int rleft, int rright) { long long ret = 0; if (A_root < B_root) { std::swap(A_root, B_root); std::swap(A, B); } // A_root >= B_root int i = A_root; while (B_root < i) { i = A.l[i]; if (A.r[i] != -1) ret += stright_right(A, A.r[i]) + A.s[A.r[i]]; } ret += A_root - B_root; int A_left = i, A_right = A_root; int B_left = B_root, B_right = B_root; while (A_left != B_left || A_right != B_right) { if (A_left != B_left) { Tree* Y; int* X_left; int* Y_left; if (A_left < B_left) { Y = &B; X_left = &A_left; Y_left = &B_left; } else { Y = &A; X_left = &B_left; Y_left = &A_left; } while (*X_left < *Y_left) { *Y_left = Y->l[*Y_left]; if (Y->r[*Y_left] != -1) ret += stright_right(*Y, Y->r[*Y_left]) + Y->s[Y->r[*Y_left]]; } } if (A_right != B_right) { Tree* Y; int* X_right; int* Y_right; if (A_right > B_right) { Y = &B; X_right = &A_right; Y_right = &B_right; } else { Y = &A; X_right = &B_right; Y_right = &A_right; } while (*X_right > *Y_right) { *Y_right = Y->r[*Y_right]; if (Y->l[*Y_right] != -1) ret += stright_left(*Y, Y->l[*Y_right]) + Y->s[Y->l[*Y_right]]; } } } if (A_left != rleft) ret += reroot(A.l[A_left], B.l[A_left], rleft, A_left - 1); if (A_right != rright) ret += reroot(A.r[A_right], B.r[A_right], A_right + 1, rright); return ret % MOD; } int main() { std::ios::sync_with_stdio(false); std::cin >> n; int A_root = A.read(n); int B_root = B.read(n); std::cout << reroot(A_root, B_root, 0, n - 1); return 0; }
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | #include <iostream> #include <vector> #define MOD 1000000007 struct Tree { std::vector<int> p, l, r, s; int read(int n) { p.resize(n); l.resize(n, -1); r.resize(n, -1); s.resize(n); int root = -1; for (int i = 0; i < n; ++i) { std::cin >> p[i]; --p[i]; if (p[i] < 0) { p[i] = -1; root = i; }else { if (i < p[i]) l[p[i]] = i; else r[p[i]] = i; } } calc_size(root); return root; } int calc_size(int i) { int size = 1; if (r[i] != -1) size += calc_size(r[i]); if (l[i] != -1) size += calc_size(l[i]); s[i] = size; return size; } }; int n; Tree A, B; long long stright_left(Tree const& T, int node); long long stright_right(Tree const& T, int node) { long long ret = 0; if (T.r[node] != -1) { ret += stright_right(T, T.r[node]); } if (T.l[node] != -1) { ret += stright_left(T, T.l[node]) + T.s[T.l[node]]; } return ret % MOD; } long long stright_left(Tree const& T, int node) { long long ret = 0; if (T.r[node] != -1) { ret += stright_right(T, T.r[node]) + T.s[T.r[node]]; } if (T.l[node] != -1) { ret += stright_left(T, T.l[node]); } return ret % MOD; } long long reroot(int A_root, int B_root, int rleft, int rright) { long long ret = 0; if (A_root < B_root) { std::swap(A_root, B_root); std::swap(A, B); } // A_root >= B_root int i = A_root; while (B_root < i) { i = A.l[i]; if (A.r[i] != -1) ret += stright_right(A, A.r[i]) + A.s[A.r[i]]; } ret += A_root - B_root; int A_left = i, A_right = A_root; int B_left = B_root, B_right = B_root; while (A_left != B_left || A_right != B_right) { if (A_left != B_left) { Tree* Y; int* X_left; int* Y_left; if (A_left < B_left) { Y = &B; X_left = &A_left; Y_left = &B_left; } else { Y = &A; X_left = &B_left; Y_left = &A_left; } while (*X_left < *Y_left) { *Y_left = Y->l[*Y_left]; if (Y->r[*Y_left] != -1) ret += stright_right(*Y, Y->r[*Y_left]) + Y->s[Y->r[*Y_left]]; } } if (A_right != B_right) { Tree* Y; int* X_right; int* Y_right; if (A_right > B_right) { Y = &B; X_right = &A_right; Y_right = &B_right; } else { Y = &A; X_right = &B_right; Y_right = &A_right; } while (*X_right > *Y_right) { *Y_right = Y->r[*Y_right]; if (Y->l[*Y_right] != -1) ret += stright_left(*Y, Y->l[*Y_right]) + Y->s[Y->l[*Y_right]]; } } } if (A_left != rleft) ret += reroot(A.l[A_left], B.l[A_left], rleft, A_left - 1); if (A_right != rright) ret += reroot(A.r[A_right], B.r[A_right], A_right + 1, rright); return ret % MOD; } int main() { std::ios::sync_with_stdio(false); std::cin >> n; int A_root = A.read(n); int B_root = B.read(n); std::cout << reroot(A_root, B_root, 0, n - 1); return 0; } |