#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; } |
English