test/DirectedMinimumSpanningTree.test.cpp

• View this file on GitHub
• Last update: 2021-03-22 17:28:05+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/all/GRL_2_B

Depends on

• Graph/CycleDetection.cpp
• Graph/DirectedMinimumSpanningTree.cpp
• Graph/GraphTemplate.cpp

Code

#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/all/GRL_2_B"
#include "./../Graph/DirectedMinimumSpanningTree.cpp"
#include <iostream>
using namespace std;

int main() {
	cin.tie(nullptr);
	ios_base::sync_with_stdio(false);

	int n, m, r;
	cin >> n >> m >> r;
	DirectedMinimumSpanningTree dmst(n);
	for (int i = 0; i < m; ++i) {
		int s, t;
		Weight cost;
		cin >> s >> t >> cost;
		dmst.add_edge(s, t, cost);
	}
	cout << dmst.solve(r) << '\n';
}

#line 1 "test/DirectedMinimumSpanningTree.test.cpp"
#define PROBLEM "https://onlinejudge.u-aizu.ac.jp/courses/library/5/GRL/all/GRL_2_B"
#line 2 "Graph/GraphTemplate.cpp"
#include <vector>
#include <utility>
#include <iostream>
#include <limits>

using Weight = long long;
constexpr Weight INF = std::numeric_limits<Weight>::max();
struct Edge {
	int to;
	Weight cost;
	Edge() : to(-1), cost(-1) {}
	Edge(int _to, Weight _cost = 1) : to(_to), cost(_cost) {}
	friend bool operator<(const Edge& e1, const Edge& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge& e1, const Edge& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge& e) {
		return os << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedGraph = std::vector<std::vector<int>>;
using Graph = std::vector<std::vector<Edge>>;
struct Edge2 {
	int from, to;
	Weight cost;
	Edge2() : from(-1), to(-1), cost(0) {}
	Edge2(int _from, int _to, Weight _cost) : from(_from), to(_to), cost(_cost) {}
	friend bool operator<(const Edge2& e1, const Edge2& e2) {
		return e1.cost < e2.cost;
	}
	friend bool operator>(const Edge2& e1, const Edge2& e2) {
		return e1.cost > e2.cost;
	}
	friend std::ostream& operator<<(std::ostream& os, const Edge2& e) {
		return os << e.from << "->" << e.to << '(' << e.cost << ')';
	}
};
using UnWeightedEdges = std::vector<std::pair<int, int>>;
using Edges = std::vector<Edge2>;
using Matrix = std::vector<std::vector<Weight>>;

auto add_edge(UnWeightedGraph& graph, int v, int u) {
	graph[v].push_back(u);
	graph[u].push_back(v);
}
auto add_edge(Graph& graph, int v, int u, Weight cost) {
	graph[v].emplace_back(u, cost);
	graph[u].emplace_back(v, cost);
}
auto to_graph(const UnWeightedGraph& graph, Weight cost = 1) {
	Graph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			result[i].emplace_back(v, cost);
		}
	}
	return result;
}
auto to_unweighted_graph(const Graph& graph) {
	UnWeightedGraph result(graph.size());
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			result[i].push_back(v);
		}
	}
	return result;
}
auto to_edges(const UnWeightedGraph& graph, bool unique = false) {
	std::vector<std::pair<int, int>> edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (int v : graph[i]) {
			if (!unique || static_cast<int>(i) < v) edges.emplace_back(i, v);
		}
	}
	return edges;
}
auto to_edges(const Graph& graph) {
	Edges edges;
	for (std::size_t i = 0; i < graph.size(); ++i) {
		for (auto [v, cost] : graph[i]) {
			edges.emplace_back(i, v, cost);
		}
	}
	return edges;
}
#line 3 "Graph/CycleDetection.cpp"
#include <stack>
#include <algorithm>
#include <optional>

std::optional<std::vector<int>> CycleDetectionEdge(
    int n, const std::vector<std::pair<int, int>>& edges, bool directed = true) {
	std::vector<std::vector<std::pair<int, int>>> graph(n);
	for (std::size_t i = 0; i < edges.size(); ++i) {
		auto [u, v] = edges[i];
		graph[u].emplace_back(v, i);
		if (!directed) graph[v].emplace_back(u, i);
	}
	std::vector<bool> visited(n), finished(n);
	std::stack<std::pair<int, int>> st;
	std::optional<int> pos;
	auto dfs = [&](auto&& f, int v, int prev_edge) -> void {
		visited[v] = true;
		for (auto [u, i] : graph[v]) {
			if (!finished[u] && (directed || i != prev_edge)) {
				st.emplace(i, v);
				if (visited[u]) {
					pos = u;
					return;
				}
				f(f, u, i);
				if (pos) return;
				st.pop();
			}
		}
		finished[v] = true;
	};
	for (int i = 0; i < n && !pos; ++i) {
		if (!visited[i]) dfs(dfs, i, -1);
	}

	if (pos) {
		std::vector<int> cycle;
		while (!st.empty()) {
			auto [top, from] = st.top();
			st.pop();
			cycle.push_back(top);
			if (from == *pos) {
				break;
			}
		}
		std::reverse(cycle.begin(), cycle.end());
		return cycle;
	} else {
		return std::nullopt;
	}
}

std::optional<std::vector<int>> CycleDetectionVertex(
    const std::vector<std::vector<int>>& graph, bool directed = true) {
	int n = graph.size();
	std::vector<bool> visited(n), finished(n);
	std::stack<int> st;
	std::optional<int> pos;
	auto dfs = [&](auto&& f, int v, int prev) -> void {
		visited[v] = true;
		st.push(v);
		for (int u : graph[v]) {
			if (!finished[u] && (directed || u != prev)) {
				if (visited[u]) {
					pos = u;
					return;
				}
				f(f, u, v);
				if (pos) return;
			}
		}
		finished[v] = true;
		st.pop();
	};
	for (int i = 0; i < n && !pos; ++i) {
		if (!visited[i]) dfs(dfs, i, -1);
	}

	if (pos) {
		std::vector<int> cycle;
		while (!st.empty()) {
			int top = st.top();
			st.pop();
			cycle.push_back(top);
			if (top == *pos) {
				break;
			}
		}
		std::reverse(cycle.begin(), cycle.end());
		return cycle;
	} else {
		return std::nullopt;
	}
}

auto CycleDetectionVertex(int n, const std::vector<std::pair<int, int>>& edges,
                          bool directed = true) {
	std::vector<std::vector<int>> graph(n);
	for (auto [u, v] : edges) {
		graph[u].push_back(v);
		if (!directed) graph[v].push_back(u);
	}
	return CycleDetectionVertex(graph, directed);
}
#line 6 "Graph/DirectedMinimumSpanningTree.cpp"
#include <numeric>

class DirectedMinimumSpanningTree {
	int n;
	Edges edges;

public:
	DirectedMinimumSpanningTree(int _n) : n(_n) {}
	void add_edge(int from, int to, Weight cost) {
		if (from != to) {
			edges.emplace_back(from, to, cost);
		}
	}
	Weight solve(int root) {
		std::vector<std::vector<Edge>> in_edges(n);
		for (auto edge : edges) {
			if (root != edge.to) {
				in_edges[edge.to].emplace_back(edge.from, edge.cost);
			}
		}

		std::vector<std::pair<int, int>> min_edges;
		std::vector<Weight> min_costs;
		for (int i = 0; i < n; ++i) {
			if (in_edges[i].empty()) continue;
			std::pair<int, int> min_edge;
			Weight min_cost = INF;
			for (auto edge : in_edges[i]) {
				if (min_cost > edge.cost) {
					min_edge = std::pair(edge.to, i);
					min_cost = edge.cost;
				}
			}
			min_edges.push_back(min_edge);
			min_costs.push_back(min_cost);
		}

		if (auto cycle_opt = CycleDetectionEdge(n, min_edges); cycle_opt) {
			const auto& cycle = cycle_opt.value();

			std::vector<Weight> minus(n);
			std::vector<bool> on_cycle(n);
			Weight cycle_cost_sum = 0;
			for (int i : cycle) {
				minus[min_edges[i].second] = min_costs[i];
				on_cycle[min_edges[i].first] = true;
				cycle_cost_sum += min_costs[i];
			}

			int compressed_vertex = min_edges[cycle.front()].first;
			DirectedMinimumSpanningTree sub(n);
			for (std::size_t i = 0; i < edges.size(); ++i) {
				auto edge = edges[i];
				if (on_cycle[edge.to]) {
					if (on_cycle[edge.from]) continue;
					edge.cost -= minus[edge.to];
					edge.to = compressed_vertex;
				}
				if (on_cycle[edge.from]) {
					edge.from = compressed_vertex;
				}
				sub.add_edge(edge.from, edge.to, edge.cost);
			}
			return cycle_cost_sum + sub.solve(root);
		} else {
			return std::accumulate(min_costs.begin(), min_costs.end(), Weight(0));
		}
	}
};
#line 4 "test/DirectedMinimumSpanningTree.test.cpp"
using namespace std;

int main() {
	cin.tie(nullptr);
	ios_base::sync_with_stdio(false);

	int n, m, r;
	cin >> n >> m >> r;
	DirectedMinimumSpanningTree dmst(n);
	for (int i = 0; i < m; ++i) {
		int s, t;
		Weight cost;
		cin >> s >> t >> cost;
		dmst.add_edge(s, t, cost);
	}
	cout << dmst.solve(r) << '\n';
}

Back to top page