competition-library/class_Chuliu_Edmonds.cpp at master · mdj982/competition-library · GitHub

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
#include "auto_util_header.hpp"

class Chuliu_Edmonds {
private:
	struct edge_t {
		int id;
		ll cost;
		vi included_stk;
		bool operator<(const edge_t & another) const {
			return cost > another.cost;
		}
	};
	// edges are directed to the node itself
	struct node {
		int overnode; bool done; bool fin; priority_queue<edge_t> edges; edge_t from;
	};
	vector<node> nodes;
	int n, root;
	vi stk;
	bool no_mca;
	int topnode(int k) {
		int a = k;
		while (nodes[a].overnode != -1) {
			a = nodes[a].overnode;
		}
		if (k != a) nodes[k].overnode = a;
		return a;
	}
	void contract(int s) {
		int a = s;
		priority_queue<edge_t> new_from_edges;
		int cnt = 0;
		do {
			a = topnode(a);
			while (nodes[a].edges.size()) {
				edge_t edge = nodes[a].edges.top();
				nodes[a].edges.pop();
				if (edge.id == nodes[a].from.id) continue;
				edge.cost -= nodes[a].from.cost;
				edge.included_stk.push_back(a);
				new_from_edges.push(edge);
			}
			nodes[a].overnode = nodes.size();
			a = nodes[a].from.id;
		} while (a != s);
		nodes.push_back({ -1, false, false, new_from_edges,{} });
	}
	void unfold() {
		while (stk.size()) {
			int a = stk.back(); stk.pop_back();
			if (a >= n) {
				int b = nodes[a].from.included_stk.back();
				ll d = nodes[b].from.cost;
				nodes[b].from = nodes[a].from;
				nodes[b].from.cost += d;
				nodes[b].from.included_stk.pop_back();
			}
			else nodes[a].fin = true;
		}
	}
public:
	Chuliu_Edmonds(const vvi &lst, const vvll &cst, int start) {
		n = lst.size();
		nodes.resize(n);
		Loop(i, n) nodes[i] = { -1, false, false, priority_queue<edge_t>(),{} };
		Loop(i, n) {
			Loop(j, lst[i].size()) {
				nodes[lst[i][j]].edges.push({ i, cst[i][j], {} });
			}
		}
		root = start;
		no_mca = false;
		nodes[root].fin = nodes[root].done = true;
		Loop(i, n) {
			if (!nodes[i].fin) {
				int a = i;
				nodes[a].done = true;
				stk.push_back(a);
				do {
					int b;
					do {
						if (nodes[a].edges.empty()) { no_mca = true; return; }
						nodes[a].from = nodes[a].edges.top(); nodes[a].edges.pop();
						b = nodes[a].from.id;
					} while (topnode(a) == topnode(b));
					if (nodes[b].fin) unfold();
					else if (nodes[b].done) {
						contract(b);
						stk.push_back(nodes.size() - 1);
						a = nodes.size() - 1;
					}
					else {
						nodes[b].done = true;
						stk.push_back(b);
						a = b;
					}
				} while (stk.size());
			}
		}
		return;
	}
	vector<P> get_tree_idpair() {
		if (no_mca) return{};
		vector<P> ret;
		Loop(i, n) {
			if (i != root) ret.push_back({ nodes[i].from.id, i });
		}
		return ret;
	}
	ll get_weight() {
		if (no_mca) return -1;
		ll ret = 0;
		Loop(i, n) {
			if (i != root) ret += nodes[i].from.cost;
		}
		return ret;
	}
};