Optimizing Prisoner Allocation Using Union-Find and Binary Search

The problem involves distributing N prisoners into two separate prisons based on M pairs of conflicts. Each conflict pair is defined by two prisoner IDs and a conflict weight. The objective is to arrange the prisoners such that the maximum conflict weight among any two prisoners sharing the same prison is minimized. We need to determine this minimum possible maximum conflict value.

Approach 1: Weighted Union-Find with Greedy Strategy

This method treats the problem similarly to constructing a maximum spanning tree or resolving constraints in a specific order. The core idea is to process conflicts in descending order of weight. We attempt to separate prisoners who have high conflict values first.

We utilize a Disjoint Set Union (DSU) data structure, often implemented with path compression and union by rank. However, since we need to track two opposing groups (Prison A and Prison B), we use a weighted DSU approach. In this implementation, for each prisoner i, i + N represents that prisoner in the opposing group. If prisoner x is in Prison A, then x + N implies x is in Prison B.

As we iterate through the edges sorted by weight:

  1. Check for Conflict: If prisoners u and v are currently in the same set, it means they are assigned to the same prison. Since we are processing edges from highest to lowest weight, this confirms that we cannot separate them without violating a higher or equal weight constraint. Thus, the current weight is the answer.
  2. Separate Groups: If they are in different sets, we must ensure they are in different prisons. We union u with the opposite of v (v + offset), and v with the opposite of u (u + offset).

If the loop completes without finding any pair in the same set, it is possible to separate all conflicting pairs, and the answer is 0.

#include <iostream>
#include <vector>
#include <algorithm>

struct Relation {
    int first, second, weight;
};

class DisjointSet {
    std::vector<int> parent, rank_;
public:
    DisjointSet(int size) : parent(size * 2), rank_(size * 2, 0) {
        for (int i = 0; i < size * 2; ++i) parent[i] = i;
    }

    int findRoot(int x) {
        if (parent[x] != x) {
            parent[x] = findRoot(parent[x]);
        }
        return parent[x];
    }

    void unionSets(int x, int y) {
        int xRoot = findRoot(x);
        int yRoot = findRoot(y);
        if (xRoot == yRoot) return;

        if (rank_[xRoot] < rank_[yRoot]) {
            parent[xRoot] = yRoot;
        } else if (rank_[xRoot] > rank_[yRoot]) {
            parent[yRoot] = xRoot;
        } else {
            parent[yRoot] = xRoot;
            rank_[xRoot]++;
        }
    }

    bool areConnected(int x, int y) {
        return findRoot(x) == findRoot(y);
    }
};

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

    int numPrisoners, numRelations;
    if (!(std::cin >> numPrisoners >> numRelations)) return 0;

    std::vector<Relation> data(numRelations);
    for (int i = 0; i < numRelations; ++i) {
        std::cin >> data[i].first >> data[i].second >> data[i].weight;
        // Converting to 0-based index
        data[i].first--;
        data[i].second--;
    }

    // Sort by weight in descending order
    std::sort(data.begin(), data.end(), [](const Relation& a, const Relation& b) {
        return a.weight > b.weight;
    });

    DisjointSet dsu(numPrisoners);
    int answer = 0;

    for (const auto& rel : data) {
        if (dsu.areConnected(rel.first, rel.second)) {
            answer = rel.weight;
            break;
        }
        // Place them in separate prisons
        dsu.unionSets(rel.first, rel.second + numPrisoners);
        dsu.unionSets(rel.second, rel.first + numPrisoners);
    }

    std::cout << answer << std::endl;
    return 0;
}

Approach 2: Binary Search Combined with Bipartite Checking

An alternative solution involves binary searching the answer. We determine if it is possible to arrange prisoners such that the maximum conflict weight inside any prison is less than or equal to a value mid.

For a given mid, we construct a graph containing only the edges where the conflict weight w is strictly greater than mid. These edges represent relationships that must be separated (i.e., the two prisoners must be in different prisons). If this graph is bipartite, then a valid assignment exists for this mid. If the graph contains an odd-length cycle (making it non-bipartite), it is impossible to separate all required conflicting pairs, meaning mid is too low.

We perform a binary search on the answer range. Since weights can be large, the range typically spans from 0 to the maximum possible weight (e.g., 10^9).

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>

using namespace std;

struct Edge {
    int u, v, cost;
};

bool checkBipartite(int n, const vector<vector<int>>& graph) {
    vector<int> color(n, 0);
    for (int start = 0; start < n; ++start) {
        if (color[start] != 0) continue;
        
        // BFS/DFS to 2-color the component
        function<bool(int, int)> dfs = [&](int node, int currentColor) {
            color[node] = currentColor;
            for (int neighbor : graph[node]) {
                if (color[neighbor] == currentColor) {
                    return false; // Conflict found
                }
                if (color[neighbor] == 0) {
                    if (!dfs(neighbor, -currentColor)) {
                        return false;
                    }
                }
            }
            return true;
        };
        
        if (!dfs(start, 1)) {
            return false;
        }
    }
    return true;
}

void solve() {
    int prisoners, pairs;
    if (!(cin >> prisoners >> pairs)) return;

    vector<Edge> edges(pairs);
    for (int i = 0; i < pairs; ++i) {
        cin >> edges[i].u >> edges[i].v >> edges[i].cost;
        edges[i].u--;
        edges[i].v--;
    }

    // Sorting is optional for the binary search logic but helps in pruning early in check
    sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b) {
        return a.cost > b.cost;
    });

    auto isValid = [&](int threshold) {
        vector<vector<int>> adjacency(prisoners);
        for (const auto& e : edges) {
            if (e.cost <= threshold) break;
            adjacency[e.u].push_back(e.v);
            adjacency[e.v].push_back(e.u);
        }
        return checkBipartite(prisoners, adjacency);
    };

    int low = 0, high = 1e9;
    while (low < high) {
        int mid = low + (high - low) / 2;
        if (isValid(mid)) {
            high = mid;
        } else {
            low = mid + 1;
        }
    }
    cout << low << "\n";
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    solve();
    return 0;
}

Tags: C++ algorithms graph theory Union-Find Binary Search

Posted on Fri, 03 Jul 2026 16:29:09 +0000 by PHPSpirit