|
| 1 | +#pragma once |
| 2 | +#include <algorithm> |
| 3 | +#include <cassert> |
| 4 | +#include <numeric> |
| 5 | +#include <utility> |
| 6 | +#include <vector> |
| 7 | + |
| 8 | +// Max size min weight two spanning forests |
| 9 | +// Complexity: O(NM + M log M) |
| 10 | +// Reference: https://www.slideshare.net/tmaehara/ss-17402143 |
| 11 | +// Verified: |
| 12 | +// - https://www.codechef.com/submit/HAMEL |
| 13 | +// - https://community.topcoder.com/stat?c=problem_statement&pm=14909&rd=17198 |
| 14 | +template <class Label, class W = int> |
| 15 | +std::pair<std::vector<bool>, std::vector<bool>> |
| 16 | +edge_disjoint_min_spanning_forests(const std::vector<std::pair<Label, Label>> &edges, |
| 17 | + const std::vector<W> &ws = {}) { |
| 18 | + assert(ws.empty() or ws.size() == edges.size()); |
| 19 | + const int M = edges.size(); |
| 20 | + std::vector<Label> lbl(M * 2); |
| 21 | + for (int e = 0; e < M; e++) lbl[e * 2] = edges[e].first, lbl[e * 2 + 1] = edges[e].second; |
| 22 | + std::sort(lbl.begin(), lbl.end()); |
| 23 | + lbl.erase(std::unique(lbl.begin(), lbl.end()), lbl.end()); |
| 24 | + const int N = lbl.size(); |
| 25 | + |
| 26 | + std::vector<std::pair<int, int>> uvs(M); |
| 27 | + for (int e = 0; e < M; e++) { |
| 28 | + int u = std::lower_bound(lbl.begin(), lbl.end(), edges[e].first) - lbl.begin(); |
| 29 | + int v = std::lower_bound(lbl.begin(), lbl.end(), edges[e].second) - lbl.begin(); |
| 30 | + uvs[e] = {u, v}; |
| 31 | + } |
| 32 | + |
| 33 | + std::vector<int> es(M); |
| 34 | + std::iota(es.begin(), es.end(), 0); |
| 35 | + if (!ws.empty()) std::sort(es.begin(), es.end(), [&](int i, int j) { return ws[i] < ws[j]; }); |
| 36 | + |
| 37 | + std::vector<std::vector<bool>> I(2, std::vector<bool>(M)); |
| 38 | + std::vector<std::vector<std::pair<int, int>>> to(N); |
| 39 | + |
| 40 | + int nb_accepted_edges = 0; |
| 41 | + auto accept_edge = [&](int e) { |
| 42 | + nb_accepted_edges++; |
| 43 | + int u = uvs[e].first, v = uvs[e].second; |
| 44 | + to[u].emplace_back(v, e); |
| 45 | + to[v].emplace_back(u, e); |
| 46 | + }; |
| 47 | + |
| 48 | + auto dfs = [&](int head, const std::vector<bool> &I) -> std::vector<std::pair<int, int>> { |
| 49 | + std::vector<int> st{head}; |
| 50 | + std::vector<std::pair<int, int>> prv(N, {-1, -1}); |
| 51 | + prv[head] = {head, -1}; |
| 52 | + while (!st.empty()) { |
| 53 | + int now = st.back(); |
| 54 | + st.pop_back(); |
| 55 | + for (auto p : to[now]) { |
| 56 | + int nxt = p.first, e = p.second; |
| 57 | + if (!I[e] or prv[nxt].first >= 0) continue; |
| 58 | + prv[nxt] = {now, e}, st.push_back(nxt); |
| 59 | + } |
| 60 | + } |
| 61 | + return prv; |
| 62 | + }; |
| 63 | + |
| 64 | + std::vector<int> prveid(M, -1), visited(N); |
| 65 | + std::vector<std::vector<int>> L(2); |
| 66 | + std::vector<std::vector<std::pair<int, int>>> prv(2); |
| 67 | + for (const int e : es) { |
| 68 | + if (nb_accepted_edges > 2 * (N - 1)) break; |
| 69 | + const int u = uvs[e].first, v = uvs[e].second; |
| 70 | + |
| 71 | + bool found = false; |
| 72 | + |
| 73 | + for (int d = 0; d < 2; d++) { |
| 74 | + prv[d] = dfs(u, I[d]); |
| 75 | + if (prv[d][v].first < 0) { |
| 76 | + accept_edge(e); |
| 77 | + I[d][e] = 1; |
| 78 | + found = true; |
| 79 | + break; |
| 80 | + } |
| 81 | + } |
| 82 | + if (found) continue; |
| 83 | + |
| 84 | + visited.assign(N, 0); |
| 85 | + visited[u] = 1; |
| 86 | + L[0] = {e}, L[1] = {}; |
| 87 | + |
| 88 | + int ehead = -1; |
| 89 | + prveid[e] = -1; |
| 90 | + for (int i = 0;; i ^= 1) { |
| 91 | + if (L[i].empty()) break; |
| 92 | + L[i ^ 1].clear(); |
| 93 | + while (!L[i].empty()) { |
| 94 | + const int exy = L[i].back(); |
| 95 | + L[i].pop_back(); |
| 96 | + int x = uvs[exy].first, y = uvs[exy].second; |
| 97 | + if (prv[i][x].first < 0 or prv[i][y].first < 0) { |
| 98 | + ehead = exy; |
| 99 | + break; |
| 100 | + } |
| 101 | + if (!visited[x]) std::swap(x, y); |
| 102 | + while (!visited[y]) { |
| 103 | + int nxty = prv[i][y].first, nxte = prv[i][y].second; |
| 104 | + L[i ^ 1].push_back(nxte); |
| 105 | + visited[y] = 1; |
| 106 | + y = nxty; |
| 107 | + prveid[nxte] = exy; |
| 108 | + } |
| 109 | + } |
| 110 | + if (ehead >= 0) { |
| 111 | + accept_edge(e); |
| 112 | + int c = I[0][ehead]; |
| 113 | + for (; ehead >= 0; ehead = prveid[ehead], c ^= 1) { |
| 114 | + I[c][ehead] = 1, I[c ^ 1][ehead] = 0; |
| 115 | + } |
| 116 | + break; |
| 117 | + } |
| 118 | + } |
| 119 | + } |
| 120 | + return {I[0], I[1]}; |
| 121 | +} |
0 commit comments