Find max cost edge between two nodes in tree - c++

Question

Task:
Given a weighted tree graph and a set of nodes pairs. For each pair (u,v) from set I need to find(effectively) maximum edge between (u,v).
My approach:
Using Tarjan's algoritm for each pair (u,v) we can find lowest common ancestor LCA(u,v) = a. Then we can present the path between (u,v) as union of (u,a) and (v,a) pathes and maximum edge between (u,v) as max(max_edge(u,a),max_edge(v,a)).
Problem:
I'm trying to add max_edge save in LCA algoritm, but haven't got any success yet.
The question is: How can I add support of max edge save at LCA Tarjan's algorithm?
My attempt code:
int max_cost;
int dsu_find(int node)
{
if (node == parent[node])
return node;
max_cost = std::max(max_cost, edges[node][parent[node]]);
return parent[node] = dsu_find(parent[node]);
}
void lca_dfs(int node, std::vector<std::list<int>> &query_list)
{
dsu_make(node);
ancestor[node] = node;
marks[node] = true;
for(auto neighbour:adjacency_list[node])
{
if (!marks[neighbour.first])
{
lca_dfs(neighbour.first,query_list);
dsu_unite(node, neighbour.first);
ancestor[dsu_find(node)] = node;
}
}
for (auto query_node : query_list[node])
if (marks[query_node])
{
dsu_find(query_node);
dsu_find(node);
printf("%d %d -> %lld\n", node, query_node,max_cost);
query_list[query_node].remove(node);
max_cost = 0;
}
}
But it's working incorrect.
My full lca implementation(without incorrect modifications):
std::vector<int> parent;
std::vector<int> rank;
std::vector<int> ancestor;
std::vector<bool> marks;
std::vector<std::list<std::pair<int, long long>>> adjacency_list;
void lca_dfs(int node, std::vector<std::list<int>> &query_list)
{
dsu_make(node);
ancestor[node] = node;
marks[node] = true;
for(auto neighbour:adjacency_list[node])
{
if (!marks[neighbour.first])
{
lca_dfs(neighbour.first,query_list);
dsu_unite(node, neighbour.first);
ancestor[dsu_find(node)] = node;
}
}
for (auto query_node : query_list[node])
if (marks[query_node])
{
printf("LCA of %d %d is %d\n", node, query_node,ancestor[dsu_find(query_node)]);
query_list[query_node].remove(node);
}
}
//dsu operations
void dsu_make(int node)
{
parent[node] = node;
rank[node] = 0;
}
int dsu_find(int node)
{
return node == parent[node] ? node : parent[node]=dsu_find(parent[node]);
}
void dsu_unite(int node_1,int node_2)
{
int root_1 = dsu_find(node_1), root_2 = dsu_find(node_2);
if(root_1!=root_2)
{
if(rank[root_1] < rank[root_2])
std::swap(root_1, root_2);
parent[root_2] = root_1;
if (rank[root_1] == rank[root_2])
rank[root_1]++;
}
}
*For each node query_list[node] consists of v such as (node,v) is needed pair.
I understood, that I using double memory(just for easier access).
I'll be gratefull for any hints or implementation fixes.

Answer 1

Hope this Implementation works for you.
#include <bits/stdc++.h>
#include <algorithm>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define ass 1e18
#define MOD 1000000007
#define mp make_pair
#define pb push_back
#define pf push_front
#define pob pop_back
#define pof pop_front
#define fi first
#define se second
#define sz(x) (ll)x.size()
#define present(c,x) ((c).find(x) != (c).end())
#define boost ios_base::sync_with_stdio(false);cin.tie(NULL);
#define debug(x) cout << #x << ": " << x << endl;
#define debug2(x,y) cout<<#x<<": "<< x<< ", "<< #y<< ": "<< y<< endl;
#define debug3(x,y,z) cout<<#x<<": "<< x<< ", "<< #y<< ": "<< y<<" "<<#z<<" : "<<z<< endl;
using namespace std;
typedef long long int ll;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ordered_set tree<ll, null_type,less<ll>, rb_tree_tag,tree_order_statistics_node_update>
pair<int,int> parent[200005],par[200005],dpp[200005][19];
vector<pair< pair<int,int>,pair<int,int> > >v;
vector<pair<int,int> >vv[200005];
int level[200005],vis[200005];
ll ans[200005];
void dfs(int x,int p)
{
level[x]=level[p]+1;
for(int i=0;i<sz(vv[x]);i++)
{
if(vv[x][i].fi!=p)
{
par[vv[x][i].fi].fi=x;
par[vv[x][i].fi].se=vv[x][i].se;
dfs(vv[x][i].fi,x);
}
}
}
void computeparent(int n)
{
for(int i=1;i<=n;i++)
dpp[i][0]=par[i];
for(int j=1;j<=18;j++)
{
for(int i=1;i<=n;i++)
{
dpp[i][j].fi=dpp[dpp[i][j-1].fi][j-1].fi;
dpp[i][j].se=max(dpp[i][j-1].se,dpp[dpp[i][j-1].fi][j-1].se);
}
}
}
int lca(int a,int b)
{
if(level[b]>level[a])
swap(a,b);
int diff=level[a]-level[b];
int m=ceil(log2(diff));
for(int i=m;i>=0;i--)
{
if(diff&(1LL<<i))
a=dpp[a][i].fi;
}
if(a==b)
return a;
for(int i=m;i>=0;i--)
{
if(dpp[a][i].fi!=dpp[b][i].fi)
{
a=dpp[a][i].fi;
b=dpp[b][i].fi;
}
}
return dpp[a][0].fi;
}
int lca2(int a,int b)
{
int c=a,d=b;
if(level[b]>level[a])
swap(a,b);
int i,maxi=0,diff=level[a]-level[b];
int m=ceil(log2(diff));
for(int i=m;i>=0;i--)
{
if(diff&(1LL<<i))
{
maxi=max(maxi,dpp[a][i].se);
a=dpp[a][i].fi;
}
}
return maxi;
}
int finds(int a)
{
while(parent[a].fi!=a)
{
a=parent[a].fi;
}
return a;
}
void unions(int x,int y)
{
if(parent[x].se>parent[y].se)
parent[y].fi=x;
else if(parent[x].se<parent[y].se)
parent[x].fi=y;
else
{
parent[x].fi=y;
parent[x].se++;
}
}
void solve()
{
ll sum=0;
int n,m,i,a,b,c;
cin>>n>>m;
for(i=1;i<=n;i++)
{
parent[i].fi=i;
parent[i].se=0;
}
for(i=0;i<m;i++)
{
cin>>a>>b>>c;
v.pb(mp(mp(c,i),mp(a,b)));
}
sort(v.begin(),v.end());
for(i=0;i<sz(v);i++)
{
int a=v[i].se.fi,b=v[i].se.se;
int x=finds(a),y=finds(b);
if(x!=y)
{
vv[a].pb(mp(b,v[i].fi.fi));
vv[b].pb(mp(a,v[i].fi.fi));
ans[v[i].fi.se]=1;
unions(x,y);
sum+=v[i].fi.fi;
}
}
dfs(1,0);
computeparent(n);
for(i=0;i<m;i++)
{
if(ans[v[i].fi.se]==0)
{
int a=lca(v[i].se.fi,v[i].se.se);
ans[v[i].fi.se]=sum+v[i].fi.fi-max(lca2(a,v[i].se.fi),lca2(a,v[i].se.se));
}
else
ans[v[i].fi.se]=sum;
}
for(i=0;i<m;i++)
cout<<ans[i]<<"\n";
}
int main()
{
boost
int t=1;
//cin>>t;
while(t--)
{
solve();
}
return 0;
}