#include<bits/stdc++.h> using namespace std; #define REP(i, n) for(int i = 0; i < n; i++) #define FOR(i, a, b) for(int i = a; i <= b; i++) #define ST first #define ND second ostream& operator<<(ostream &out, string str) { for(char c : str) out << c; return out; } template<class L, class R> ostream& operator<<(ostream &out, pair<L, R> p) { return out << "(" << p.ST << ", " << p.ND << ")"; } template<class T> auto operator<<(ostream &out, T &&x) -> decltype(x.begin(), out) { out << '{'; for(auto &e : x) out << e << (&e == &*--x.end() ? "" : ", "); return out << '}'; } template<class... Args> void dump(Args&&... args) { ((cerr << args << "; "), ...); } #ifdef DEBUG # define debug(x...) cerr << "[" #x "]: ", dump(x), cerr << "\n" #else # define debug(...) false #endif template<class T> int size(T && a) { return (int) a.size(); } using LL = long long; using PII = pair<LL, int>; struct Node { int val; Node *l = nullptr, *r = nullptr; Node(int v) : val(v) {} ~Node() { delete l; delete r; } }; using pNode = Node*; pNode gen_tree(vector<int> &a) { int n = size(a); vector<pNode> nodes(n, nullptr); REP(i, n) nodes[i] = new Node(i); pNode root = nullptr; REP(i, n) { if(a[i] == -1) root = nodes[i]; else { int x = a[i] - 1; if(x < i) nodes[x]->r = nodes[i]; else nodes[x]->l = nodes[i]; } } return root; } PII unfold(pNode x, int side) { if(!x) return {0, 0}; auto [a, l] = unfold(x->l, 0); auto [b, r] = unfold(x->r, 1); return {a + b + (side ? l : r), l + 1 + r}; } tuple<LL, int, pNode> rem_top(pNode x, int k, int side) { if(k == 0) return {0, 0, x}; auto ops = (side ? x->l : x->r); auto dep = (side ? x->r : x->l); auto [cost, sub] = unfold(ops, side ^ 1); cost += sub, sub++; if(sub >= k) return {cost, sub - k, dep}; else { auto [cc, top, ptr] = rem_top(dep, k - sub, side); return {cost + cc, top, ptr}; } } LL dst(int top, pNode x, pNode y) { if(top == 0) { if(!x) return 0; if(x->val > y->val) swap(x, y); int dif = y->val - x->val; auto [cost, t, p] = rem_top(x->r, dif, 1); cost += dif; LL a = dst(t, p, y->r); LL b = dst(-dif, x->l, y->l); return cost + a + b; } else if(top > 0) { auto [cost, t, p] = rem_top(y, top, 1); return cost + dst(t, p, x); } else { auto [cost, t, p] = rem_top(y, -top, 0); return cost + dst(-t, p, x); } } int main() { ios_base::sync_with_stdio(0); cin.tie(0); int n; cin >> n; vector<int> a(n), b(n); REP(i, n) cin >> a[i]; REP(i, n) cin >> b[i]; auto s = gen_tree(a); auto t = gen_tree(b); cout << dst(0, s, t) << "\n"; delete s; delete t; }
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 | #include<bits/stdc++.h> using namespace std; #define REP(i, n) for(int i = 0; i < n; i++) #define FOR(i, a, b) for(int i = a; i <= b; i++) #define ST first #define ND second ostream& operator<<(ostream &out, string str) { for(char c : str) out << c; return out; } template<class L, class R> ostream& operator<<(ostream &out, pair<L, R> p) { return out << "(" << p.ST << ", " << p.ND << ")"; } template<class T> auto operator<<(ostream &out, T &&x) -> decltype(x.begin(), out) { out << '{'; for(auto &e : x) out << e << (&e == &*--x.end() ? "" : ", "); return out << '}'; } template<class... Args> void dump(Args&&... args) { ((cerr << args << "; "), ...); } #ifdef DEBUG # define debug(x...) cerr << "[" #x "]: ", dump(x), cerr << "\n" #else # define debug(...) false #endif template<class T> int size(T && a) { return (int) a.size(); } using LL = long long; using PII = pair<LL, int>; struct Node { int val; Node *l = nullptr, *r = nullptr; Node(int v) : val(v) {} ~Node() { delete l; delete r; } }; using pNode = Node*; pNode gen_tree(vector<int> &a) { int n = size(a); vector<pNode> nodes(n, nullptr); REP(i, n) nodes[i] = new Node(i); pNode root = nullptr; REP(i, n) { if(a[i] == -1) root = nodes[i]; else { int x = a[i] - 1; if(x < i) nodes[x]->r = nodes[i]; else nodes[x]->l = nodes[i]; } } return root; } PII unfold(pNode x, int side) { if(!x) return {0, 0}; auto [a, l] = unfold(x->l, 0); auto [b, r] = unfold(x->r, 1); return {a + b + (side ? l : r), l + 1 + r}; } tuple<LL, int, pNode> rem_top(pNode x, int k, int side) { if(k == 0) return {0, 0, x}; auto ops = (side ? x->l : x->r); auto dep = (side ? x->r : x->l); auto [cost, sub] = unfold(ops, side ^ 1); cost += sub, sub++; if(sub >= k) return {cost, sub - k, dep}; else { auto [cc, top, ptr] = rem_top(dep, k - sub, side); return {cost + cc, top, ptr}; } } LL dst(int top, pNode x, pNode y) { if(top == 0) { if(!x) return 0; if(x->val > y->val) swap(x, y); int dif = y->val - x->val; auto [cost, t, p] = rem_top(x->r, dif, 1); cost += dif; LL a = dst(t, p, y->r); LL b = dst(-dif, x->l, y->l); return cost + a + b; } else if(top > 0) { auto [cost, t, p] = rem_top(y, top, 1); return cost + dst(t, p, x); } else { auto [cost, t, p] = rem_top(y, -top, 0); return cost + dst(-t, p, x); } } int main() { ios_base::sync_with_stdio(0); cin.tie(0); int n; cin >> n; vector<int> a(n), b(n); REP(i, n) cin >> a[i]; REP(i, n) cin >> b[i]; auto s = gen_tree(a); auto t = gen_tree(b); cout << dst(0, s, t) << "\n"; delete s; delete t; } |