#include <iostream> #include <vector> #include <queue> #include <utility> #define PRIME 1000000007 struct Node { int value = 0; Node *left = nullptr, *right = nullptr; }; long long rollOutCost(Node *tree); int main() { std::ios_base::sync_with_stdio(false); std::cin.tie(nullptr); int n; std::cin >> n; std::vector<Node *> nodes[2]; nodes[0].resize(n + 1); nodes[1].resize(n + 1); Node *roots[2]; for (int z: {0, 1}) { for (int i = 1; i <= n; ++i) { nodes[z][i] = new Node{}; nodes[z][i]->value = i; } } for (int z: {0, 1}) { for (int i = 1; i <= n; ++i) { int iThParent; std::cin >> iThParent; if (iThParent == -1) { roots[z] = nodes[z][i]; continue; } if (nodes[z][i]->value < nodes[z][iThParent]->value) { nodes[z][iThParent]->left = nodes[z][i]; } else { nodes[z][iThParent]->right = nodes[z][i]; } } } long long numberOfRotations = 0; std::queue<std::pair<Node *, Node *>> bfsQueue; bfsQueue.push(std::make_pair(roots[0], roots[1])); while (not bfsQueue.empty()) { auto [fromNode, toNode] = bfsQueue.front(); bfsQueue.pop(); if (fromNode->value > toNode->value) std::swap(fromNode, toNode); auto fromRightIter = fromNode, toRightIter = toNode; while (fromRightIter->value != toRightIter->value) { while (fromRightIter->value < toRightIter->value) if (fromRightIter->right) fromRightIter = fromRightIter->right; else break; while (toRightIter->value < fromRightIter->value) if (toRightIter->right) toRightIter = toRightIter->right; else break; } auto fromLeftIter = fromNode, toLeftIter = toNode; while (fromLeftIter->value != toLeftIter->value) { while (fromLeftIter->value < toLeftIter->value) if (toLeftIter->left) toLeftIter = toLeftIter->left; else break; while (toLeftIter->value < fromLeftIter->value) if (fromLeftIter->left) fromLeftIter = fromLeftIter->left; else break; } if (fromLeftIter->left) { bfsQueue.push(std::make_pair(fromLeftIter->left, toLeftIter->left)); fromLeftIter->left = nullptr; toLeftIter->left = nullptr; } if (fromRightIter->right) { bfsQueue.push(std::make_pair(fromRightIter->right, toRightIter->right)); fromRightIter->right = nullptr; toRightIter->right = nullptr; } numberOfRotations += rollOutCost(fromNode); numberOfRotations += rollOutCost(toNode); numberOfRotations += toNode->value - fromNode->value; numberOfRotations %= PRIME; } std::cout << numberOfRotations << std::endl; return 0; } long long rollOutCost(Node *tree) { long long cost = 0; std::queue<Node *> nodes; nodes.push(tree); while (not nodes.empty()) { if (nodes.front()->left) { nodes.push(nodes.front()->left); cost += nodes.front()->value - nodes.front()->left->value - 1; cost %= PRIME; } if (nodes.front()->right) { nodes.push(nodes.front()->right); cost += nodes.front()->right->value - nodes.front()->value - 1; cost %= PRIME; } nodes.pop(); } return cost; }
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 | #include <iostream> #include <vector> #include <queue> #include <utility> #define PRIME 1000000007 struct Node { int value = 0; Node *left = nullptr, *right = nullptr; }; long long rollOutCost(Node *tree); int main() { std::ios_base::sync_with_stdio(false); std::cin.tie(nullptr); int n; std::cin >> n; std::vector<Node *> nodes[2]; nodes[0].resize(n + 1); nodes[1].resize(n + 1); Node *roots[2]; for (int z: {0, 1}) { for (int i = 1; i <= n; ++i) { nodes[z][i] = new Node{}; nodes[z][i]->value = i; } } for (int z: {0, 1}) { for (int i = 1; i <= n; ++i) { int iThParent; std::cin >> iThParent; if (iThParent == -1) { roots[z] = nodes[z][i]; continue; } if (nodes[z][i]->value < nodes[z][iThParent]->value) { nodes[z][iThParent]->left = nodes[z][i]; } else { nodes[z][iThParent]->right = nodes[z][i]; } } } long long numberOfRotations = 0; std::queue<std::pair<Node *, Node *>> bfsQueue; bfsQueue.push(std::make_pair(roots[0], roots[1])); while (not bfsQueue.empty()) { auto [fromNode, toNode] = bfsQueue.front(); bfsQueue.pop(); if (fromNode->value > toNode->value) std::swap(fromNode, toNode); auto fromRightIter = fromNode, toRightIter = toNode; while (fromRightIter->value != toRightIter->value) { while (fromRightIter->value < toRightIter->value) if (fromRightIter->right) fromRightIter = fromRightIter->right; else break; while (toRightIter->value < fromRightIter->value) if (toRightIter->right) toRightIter = toRightIter->right; else break; } auto fromLeftIter = fromNode, toLeftIter = toNode; while (fromLeftIter->value != toLeftIter->value) { while (fromLeftIter->value < toLeftIter->value) if (toLeftIter->left) toLeftIter = toLeftIter->left; else break; while (toLeftIter->value < fromLeftIter->value) if (fromLeftIter->left) fromLeftIter = fromLeftIter->left; else break; } if (fromLeftIter->left) { bfsQueue.push(std::make_pair(fromLeftIter->left, toLeftIter->left)); fromLeftIter->left = nullptr; toLeftIter->left = nullptr; } if (fromRightIter->right) { bfsQueue.push(std::make_pair(fromRightIter->right, toRightIter->right)); fromRightIter->right = nullptr; toRightIter->right = nullptr; } numberOfRotations += rollOutCost(fromNode); numberOfRotations += rollOutCost(toNode); numberOfRotations += toNode->value - fromNode->value; numberOfRotations %= PRIME; } std::cout << numberOfRotations << std::endl; return 0; } long long rollOutCost(Node *tree) { long long cost = 0; std::queue<Node *> nodes; nodes.push(tree); while (not nodes.empty()) { if (nodes.front()->left) { nodes.push(nodes.front()->left); cost += nodes.front()->value - nodes.front()->left->value - 1; cost %= PRIME; } if (nodes.front()->right) { nodes.push(nodes.front()->right); cost += nodes.front()->right->value - nodes.front()->value - 1; cost %= PRIME; } nodes.pop(); } return cost; } |