#include "books.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pil = pair<int, ll>;
using pli = pair<ll, int>;
const int INF = 2'000'000;
const ll LINF = 987'654'321'987'654;
struct mmset
{
int mini = INF;
int maxi = -1;
set<int> base_set;
void insert(int x)
{
base_set.insert(x);
mini = min(mini, x);
maxi = max(maxi, x);
}
int lower_bound(int x)
{
if (maxi>=x)
return *(base_set.lower_bound(x));
return INF;
}
};
vector<ll> dijkstra(int s, vector<pair<int, pil>> edgelist, int n)
{
vector<ll> ans(n, LINF);
vector<bool> processed(n, 0);
vector<vector<pil>> edges(n);
for (pair<int, pil> x: edgelist)
{
edges[x.first].push_back(x.second);
//cout << x.first << x.second.first << x.second.second << '\n';
}
ans[s] = 0;
priority_queue<pli, vector<pli>, greater<pli>> pq;
pq.push({0, s});
while (!pq.empty())
{
pli info = pq.top(); pq.pop();
ll cost = info.first;
int x = info.second;
if (processed[x]==0)
{
processed[x] = 1;
for (pil outgoing_edge: edges[x])
{
int y = outgoing_edge.first;
ll weight = outgoing_edge.second;
pq.push({cost+weight, y});
ans[y] = min(ans[y], cost+weight);
}
}
}
/*for (ll x: ans)
cout << x << ' ';
cout << '\n';*/
return ans;
}
struct ufsets
{
vector<int> uf;
vector<mmset> group_list;
ufsets(int n)
{
uf = vector<int> (n, -1);
group_list = vector<mmset> (n);
for (int i=0; i<n; i++)
group_list[i].insert(i);
}
int fin(int x)
{
if (uf[x]<0)
return x;
int y = fin(uf[x]);
uf[x] = y;
return y;
}
void mergatory(int x, int y)//add y to x
{
if (x!=y)//safety and sanity check
{
uf[x] = uf[x]+uf[y];
uf[y] = x;
for (int k: group_list[y].base_set)
group_list[x].insert(k);
group_list[y] = mmset();
}
}
void onion(int x, int y)
{
x = fin(x); y = fin(y);
if (x!=y)
{
if (uf[x]<uf[y])
mergatory(x, y);
else
mergatory(y, x);
}
}
bool contained_in_group(int x, int y)//is any member of y enclosed within x's group?
{
int lo = group_list[x].mini;
int hi = group_list[x].maxi;
if (group_list[y].lower_bound(lo)<=hi)
return 1;
return 0;
}
bool for_merge(int x, int y)
{
x = fin(x); y = fin(y);
if (contained_in_group(x, y) and contained_in_group(y, x))
{
onion(x, y);
return 1;
}
return 0;
}
void nontrivial_merges()
{
int n = uf.size();
vector<int> tree(n, -1);
vector<vector<int>> children(n);
vector<pii> current_groups;
for (int i=0; i<n; i++)
if (uf[i]<-1)//could change this to 0, but I really only care about stuff with at least two nodes in the first place
{
current_groups.push_back({group_list[i].mini, group_list[i].maxi});
}
sort(current_groups.begin(), current_groups.end());
int k = current_groups.size();
stack<pii> cycle_merger;
for (int i=0; i<k; i++)
{
pii curr_p = current_groups[i];
while ((cycle_merger.size()>0) and (curr_p.first > cycle_merger.top().second))
cycle_merger.pop();
if (cycle_merger.size()>0)
{
int parent = fin(cycle_merger.top().first);
tree[fin(curr_p.first)] = parent;
children[parent].push_back(fin(curr_p.first));
}
cycle_merger.push(curr_p);
}
//Make sure you do this top bottom, not bottom up.
queue<int> process_list;
for (int i=0; i<n; i++)
if ((tree[i]==-1) and (uf[i]<-1))//could change this to 0, but I really only care about stuff with at least two nodes in the first place
process_list.push(i);
while (process_list.size()>0)
{
int x = process_list.front();
process_list.pop();
for (int y: children[x])
{
for_merge(x, y);
process_list.push(y);
}
/*if (tree[x]!=-1)
{
for_merge(x, tree[x]);
descendant_plus_one_count[tree[x]]--;
if (descendant_plus_one_count[tree[x]]==1)
process_list.push(tree[x]);
}*/
}
}
ll get_penalty(int s)
{
int n = uf.size();
vector<int> tree(n, -1);
vector<vector<int>> children(n);
vector<pii> current_groups;
for (int i=0; i<n; i++)
if (uf[i]<-1)//could change this to 0, but I really only care about stuff with at least two nodes in the first place
{
current_groups.push_back({group_list[i].mini, group_list[i].maxi});
}
sort(current_groups.begin(), current_groups.end());
int k = current_groups.size();
stack<pii> cycle_merger;
for (int i=0; i<k; i++)
{
pii curr_p = current_groups[i];
while ((cycle_merger.size()>0) and (curr_p.first > cycle_merger.top().second))
cycle_merger.pop();
if (cycle_merger.size()>0)
{
int parent = fin(cycle_merger.top().first);
tree[fin(curr_p.first)] = parent;
children[parent].push_back(fin(curr_p.first));
}
cycle_merger.push(curr_p);
}
vector<pair<int, pil>> edgelist;
vector<int> roots;
for (int i=0; i<n; i++)
if ((tree[i]==-1) and (uf[i]<-1))//could change this to 0, but I really only care about stuff with at least two nodes in the first place
roots.push_back(i);
k = roots.size();
ll total = 0;
for (int i=0; i<(k-1); i++)
{
int a = fin(roots[i]);
int b = fin(roots[i+1]);
ll dist = group_list[b].mini - group_list[a].maxi;
edgelist.push_back({a, {b, dist}});
edgelist.push_back({b, {a, dist}});
total+=dist;
}
for (int table=0; table<n; table++)
if (children[table].size() > 0)
{
int kk = children[table].size();
for (int i=0; i<(kk-1); i++)
{
int a = fin(children[table][i]);
int b = fin(children[table][i+1]);
ll dist = group_list[b].mini - group_list[a].maxi;
edgelist.push_back({a, {b, dist}});
edgelist.push_back({b, {a, dist}});
}
int a = fin(children[table][0]);
int aa = group_list[a].mini;
while (fin(aa)!=fin(table))
{
aa--;
}
edgelist.push_back({a, {table, group_list[a].mini-aa}});
int b = fin(children[table][kk-1]);
int bb = group_list[b].maxi;
while (fin(bb)!=fin(table))
bb++;
edgelist.push_back({b, {table, bb-group_list[b].maxi}});
}
for (int i=0; i<n; i++)
{
int rep = fin(i);
edgelist.push_back({s, {rep, abs(s-i)}});
}
vector<ll> answer_key = dijkstra(s, edgelist, n);
//I can assume that there is at least one cycle to worry about
int lo_rep = roots[0];
int hi_rep = roots[k-1];
int lo = group_list[lo_rep].mini;
int hi = group_list[lo_rep].maxi;
if (s<lo)
return answer_key[hi_rep];
if (s>hi)
return answer_key[lo_rep];
ll freedom = 0;
for (int root: roots)
{
if ((group_list[root].mini <= s) and (s <= group_list[root].maxi))
freedom = answer_key[fin(root)];
}
//cout << freedom << ' ' << answer_key[lo_rep] << ' ' << answer_key[hi_rep] << '\n';
ll final_answer = total + freedom;
return final_answer;
}
};
long long minimum_walk(std::vector<int> p, int s) {
int n = p.size();
ll base_cost = 0;
for (int i=0; i<n; i++)
base_cost += abs(p[i]-i);
ufsets ufset(n);
vector<pii> cycles;
vector<bool> used(n, 0);
for (int i=0; i<n; i++)
if (used[i]==0)
{
used[i]=1;
int mini = i;
int maxi = i;
int curr = p[i];
while (curr!=i)
{
used[curr] = 1;
maxi = max(curr, maxi);
ufset.onion(i, curr);
curr = p[curr];
}
if (mini!=maxi)
cycles.push_back({mini, maxi});
}
int k = cycles.size();
if (k==0)
return 0;//this means that all books are in place
stack<pii> cycle_merger;
for (int i=0; i<k; i++)
{
pii curr_p = cycles[i];
while ((cycle_merger.size()>0) and (curr_p.first > cycle_merger.top().second))
cycle_merger.pop();
while ((cycle_merger.size()>0) and (cycle_merger.top().second < curr_p.second))
{
ufset.onion(curr_p.first, cycle_merger.top().first);
curr_p = {cycle_merger.top().first, curr_p.second};
cycle_merger.pop();
}
cycle_merger.push(curr_p);
}//the "trivial merges" should have been done by now
ufset.nontrivial_merges();
ll penalty = ufset.get_penalty(s);//determine the penalty
return base_cost + 2ll*penalty;
}
