#include <iostream> #include <stdio.h> using namespace std; const int MOD = 1000000007; const int N = 500000; int EMPTY = -1; int n; int tree[2][N + 1][4]; int PARENT = 0; int LEFT = 1; int RIGHT = 2; int SIZE = 3; int abs(int x) { return x >= 0 ? x : -x; } int linearize(int iTree, int leftBound, int rightBound) { int count = 0; for (int i = leftBound; i <= rightBound; i++) { int parent = tree[iTree][i][PARENT]; if (parent != EMPTY) { int grandparent = tree[iTree][parent][PARENT]; if (grandparent != EMPTY && leftBound <= grandparent && grandparent <= rightBound) { // if parent and grandparent are on opposite sides if (parent < i && i < grandparent || grandparent < i && i < parent) { // we need to rotate this node plus all children (=SIZE) count += tree[iTree][i][SIZE]; count %= MOD; } } } } return count; } int compare(int rootA, int rootB) { // find common left int leftA = rootA; int leftB = rootB; while (leftA != leftB) { if (leftA < leftB) { leftB = tree[1][leftB][LEFT]; } else if (leftB < leftA) { leftA = tree[0][leftA][LEFT]; } } // find common right int rightA = rootA; int rightB = rootB; while (rightA != rightB) { if (rightA < rightB) { rightA = tree[0][rightA][RIGHT]; } else if (rightB < rightA) { rightB = tree[1][rightB][RIGHT]; } } // found commmon left and right int left = leftA; int right = rightA; // include rotating to common root int result = abs(rootA - rootB); // include cost of making range linear for (int iTree = 0; iTree < 2; iTree++) { result += linearize(iTree, left, right); result %= MOD; } // include recursion to left side if (tree[0][left][LEFT] != EMPTY) { result += compare( tree[0][left][LEFT], tree[1][left][LEFT]); result %= MOD; } // include recursion to right side if (tree[0][right][RIGHT] != EMPTY) { result += compare( tree[0][right][RIGHT], tree[1][right][RIGHT]); result %= MOD; } return result; } int main() { cin >> n; int roots[2]; for (int iTree = 0; iTree < 2; iTree++) { for (int iNode = 1; iNode <= n; iNode++) { tree[iTree][iNode][LEFT] = EMPTY; tree[iTree][iNode][RIGHT] = EMPTY; } } // initialize PARENT, LEFT, RIGHT for (int iTree = 0; iTree < 2; iTree++) { for (int iNode = 1; iNode <= n; iNode++) { int parent; cin >> parent; tree[iTree][iNode][PARENT] = parent; tree[iTree][iNode][SIZE] = EMPTY; if (parent == EMPTY) { roots[iTree] = iNode; } else { if (iNode < parent) { tree[iTree][parent][LEFT] = iNode; } if (parent < iNode) { tree[iTree][parent][RIGHT] = iNode; } } } } // calculate sizes for (int iTree = 0; iTree < 2; iTree++) { int index = 1; while (true) { int left = tree[iTree][index][LEFT]; int right = tree[iTree][index][RIGHT]; if (left != EMPTY && tree[iTree][left][SIZE] == EMPTY) { index = left; continue; } if (right != EMPTY && tree[iTree][right][SIZE] == EMPTY) { index = right; continue; } int size = 1; if (left != EMPTY) { size += tree[iTree][left][SIZE]; } if (right != EMPTY) { size += tree[iTree][right][SIZE]; } tree[iTree][index][SIZE] = size; int parent = tree[iTree][index][PARENT]; if (parent == EMPTY) { break; } else { index = parent; } } } int result = compare(roots[0], roots[1]); printf("%d", result); 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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | #include <iostream> #include <stdio.h> using namespace std; const int MOD = 1000000007; const int N = 500000; int EMPTY = -1; int n; int tree[2][N + 1][4]; int PARENT = 0; int LEFT = 1; int RIGHT = 2; int SIZE = 3; int abs(int x) { return x >= 0 ? x : -x; } int linearize(int iTree, int leftBound, int rightBound) { int count = 0; for (int i = leftBound; i <= rightBound; i++) { int parent = tree[iTree][i][PARENT]; if (parent != EMPTY) { int grandparent = tree[iTree][parent][PARENT]; if (grandparent != EMPTY && leftBound <= grandparent && grandparent <= rightBound) { // if parent and grandparent are on opposite sides if (parent < i && i < grandparent || grandparent < i && i < parent) { // we need to rotate this node plus all children (=SIZE) count += tree[iTree][i][SIZE]; count %= MOD; } } } } return count; } int compare(int rootA, int rootB) { // find common left int leftA = rootA; int leftB = rootB; while (leftA != leftB) { if (leftA < leftB) { leftB = tree[1][leftB][LEFT]; } else if (leftB < leftA) { leftA = tree[0][leftA][LEFT]; } } // find common right int rightA = rootA; int rightB = rootB; while (rightA != rightB) { if (rightA < rightB) { rightA = tree[0][rightA][RIGHT]; } else if (rightB < rightA) { rightB = tree[1][rightB][RIGHT]; } } // found commmon left and right int left = leftA; int right = rightA; // include rotating to common root int result = abs(rootA - rootB); // include cost of making range linear for (int iTree = 0; iTree < 2; iTree++) { result += linearize(iTree, left, right); result %= MOD; } // include recursion to left side if (tree[0][left][LEFT] != EMPTY) { result += compare( tree[0][left][LEFT], tree[1][left][LEFT]); result %= MOD; } // include recursion to right side if (tree[0][right][RIGHT] != EMPTY) { result += compare( tree[0][right][RIGHT], tree[1][right][RIGHT]); result %= MOD; } return result; } int main() { cin >> n; int roots[2]; for (int iTree = 0; iTree < 2; iTree++) { for (int iNode = 1; iNode <= n; iNode++) { tree[iTree][iNode][LEFT] = EMPTY; tree[iTree][iNode][RIGHT] = EMPTY; } } // initialize PARENT, LEFT, RIGHT for (int iTree = 0; iTree < 2; iTree++) { for (int iNode = 1; iNode <= n; iNode++) { int parent; cin >> parent; tree[iTree][iNode][PARENT] = parent; tree[iTree][iNode][SIZE] = EMPTY; if (parent == EMPTY) { roots[iTree] = iNode; } else { if (iNode < parent) { tree[iTree][parent][LEFT] = iNode; } if (parent < iNode) { tree[iTree][parent][RIGHT] = iNode; } } } } // calculate sizes for (int iTree = 0; iTree < 2; iTree++) { int index = 1; while (true) { int left = tree[iTree][index][LEFT]; int right = tree[iTree][index][RIGHT]; if (left != EMPTY && tree[iTree][left][SIZE] == EMPTY) { index = left; continue; } if (right != EMPTY && tree[iTree][right][SIZE] == EMPTY) { index = right; continue; } int size = 1; if (left != EMPTY) { size += tree[iTree][left][SIZE]; } if (right != EMPTY) { size += tree[iTree][right][SIZE]; } tree[iTree][index][SIZE] = size; int parent = tree[iTree][index][PARENT]; if (parent == EMPTY) { break; } else { index = parent; } } } int result = compare(roots[0], roots[1]); printf("%d", result); return 0; } |