A Faster Approximation Algorithm for the Steiner Tree Problem in Graphs

steinerTreeBuilder.h

/***********************************************************************************
    From <<A Faster Approximation Algorithm for the Steiner Tree Problem in Graphs>>
        Kurt MEHLHORN

    O(N logN) 2-approximation
	
    Implement: SunMoon Master(Jinkela SHENGDIYAGE)
************************************************************************************/
#pragma once

#include<stack>
#include<deque>
#include<vector>
#include<algorithm>
#include<unordered_set>

using std::sort;
using std::vector;
using std::size_t;

class DisjointSet{
	size_t sz;
	vector<size_t> dis;
	vector<size_t> sum;
public:
	DisjointSet(size_t=0);
	void init(size_t);
	size_t find(size_t);
	bool same(size_t,size_t);
	void combine(size_t,size_t);
	size_t size(size_t);
};

DisjointSet::DisjointSet(size_t s){
	init(s);
}

void DisjointSet::init(size_t s){
	sz=s;
	dis.resize(sz);
	sum.resize(sz);
	for(size_t i=0;i<sz;++i){
		dis[i]=i;
		sum[i]=1;
	}
}

size_t DisjointSet::find(size_t p){
	if( dis[p]==p ) return p;
	return dis[p]=find(dis[p]);
}

bool DisjointSet::same(size_t a,size_t b){
	return find(a)==find(b);
}

void DisjointSet::combine(size_t a,size_t b){
	if( !same(a,b) ){
		if( sum[dis[a]] < sum[dis[b]] )
			std::swap(a,b);
		sum[dis[a]] += sum[dis[b]];
		dis[dis[b]] = dis[a];
	}
}

size_t DisjointSet::size(size_t p){
	return sum[find(p)];
}

template<class T>
struct steinerTreeBuilder{
	// 2-approximate algorithm
	// 0-base
	struct Edge{
		size_t u,v;
		T cost;
		Edge(size_t u, size_t v, const T&cost):u(u),v(v),cost(cost){}
	};
private:
	const T INF; // for shortest path algorithm relax
	vector< vector<size_t> > graph;
	vector<Edge> edge;
public:
	steinerTreeBuilder(const T &INF):INF(INF){}
	void clear(){
		graph.clear();
		edge.clear();
	}
	void init(size_t n){
		clear();
		graph.resize(n);
	}
	void add_edge(size_t u, size_t v,const T& cost){
		graph[u].emplace_back(edge.size());
		edge.emplace_back(u,v,cost);
		graph[v].emplace_back(edge.size());
		edge.emplace_back(v,u,cost);
	}
	const vector<Edge>& getEdges()const{
		return edge;
	}
	vector<size_t> solve(vector<size_t> terminalVertex){
		sort(terminalVertex.begin(),terminalVertex.end());
		terminalVertex.erase(unique(terminalVertex.begin(),terminalVertex.end()),terminalVertex.end());
		
		size_t N = graph.size();
		
		const size_t INVLID = std::numeric_limits<size_t>::max();
		
		vector<T> dist(N,INF);
		vector<size_t> prev_eid(N,INVLID);
		vector<size_t> index(N,INVLID);
		
		vector<bool> inqueue(N,false);
		std::deque<size_t> qu;
		for(size_t u:terminalVertex){
			dist[u] = 0;
			index[u]=u;
			qu.push_back(u);
			inqueue[u]=true;
		}
		
		//SPFA
		while(!qu.empty()){
			size_t u = qu.front();
			qu.pop_front();
			inqueue[u]=false;
			for(size_t eid:graph[u]){
				size_t v = edge[eid].v;
				const T &cost = edge[eid].cost;
				if( dist[v] > dist[u]+cost ){
					dist[v] = dist[u]+cost;
					prev_eid[v] = eid;
					index[v] = index[u];
					if( !inqueue[v] ){
						if( !qu.empty() && dist[v] <= dist[qu.front()] )
							qu.push_front(v);
						else
							qu.push_back(v);
						inqueue[v]=true;
					}
				}
			}
		}
		
		// Find CrossEdge
		vector< std::tuple<T,size_t> > CrossEdge;
		{
			size_t eid=0;
			for(const Edge &E:edge){
				if( index[E.u]!=index[E.v] && E.u > E.v )
					CrossEdge.emplace_back( dist[E.u]+dist[E.v]+E.cost , eid );
				eid++;
			}
		}
		sort(CrossEdge.begin(),CrossEdge.end());
		
		//Kruskal 1
		vector<size_t> SelectKEdge;
		DisjointSet ds(N);
		for(const auto &TUS:CrossEdge)
		{
			size_t eid;
			std::tie(std::ignore,eid) = TUS;
			const auto &E = edge[eid];
			if( !ds.same(index[E.u],index[E.v]) )
			{
				SelectKEdge.emplace_back(eid);
				ds.combine(index[E.u],index[E.v]);
			}
		}
		
		//Recover Edge
		vector<bool> &used = inqueue;
		used.resize(edge.size(),false);
		
		auto &FinalEdge = CrossEdge;
		FinalEdge.clear();
		for(size_t eid:SelectKEdge){
			FinalEdge.emplace_back(edge[eid].cost,eid);
			for(int t=0;t<2;++t){
				size_t v = (t&1)?edge[eid].u:edge[eid].v;
				while( v!=index[v] && !used[prev_eid[v]] ){
					used[prev_eid[v]] = true;
					FinalEdge.emplace_back(edge[prev_eid[v]].cost,prev_eid[v]);
					v = edge[prev_eid[v]].u;
				}
			}
		}
		
		//Kruskal 2
		auto &deg = dist;
		std::fill(deg.begin(),deg.end(),0);
		std::unordered_set<int> used_eid;
		
		sort(FinalEdge.begin(),FinalEdge.end());
		ds.init(N);
		SelectKEdge.clear();
		for(const auto &TUS:CrossEdge){
			size_t eid;
			std::tie(std::ignore,eid) = TUS;
			const auto &E = edge[eid];
			if( !ds.same(E.u,E.v) ){
				used_eid.insert(eid);
				deg[E.u]++;
				deg[E.v]++;
				SelectKEdge.emplace_back(eid);
				ds.combine(E.u,E.v);
			}
		}
		
		//Reduce
		vector<bool> &isTerminal = used;
		isTerminal.resize(N,false);
		for(size_t u:terminalVertex){
			isTerminal[u] = true;
		}
		std::stack< size_t, vector<size_t> > stk;
		for(size_t u=0;u<N;++u){
			if( !isTerminal[u] && deg[u]==1 ){
				deg[u]--;
				stk.emplace(u);
			}
		}
	
		while( !stk.empty() ){
			size_t u = stk.top();
			stk.pop();
	
			//find the edge to delete
			for(size_t eid:graph[u]){
				if( used_eid.find(eid) == used_eid.end() )
					continue;
				used_eid.erase(eid);
	
				auto v = edge[eid].v;
				deg[ v ]--;
				if( deg[ v ]==0 && !isTerminal[v])
					stk.emplace(v);
			}
		}
		SelectKEdge.clear();
		std::copy(used_eid.begin(),used_eid.end(),std::back_inserter(SelectKEdge));
	
		return SelectKEdge;
	}
};