13161번: 분단의 슬픔 - 최대 유량(MFMC, Dinic)

사람들을 \(A, B\) 두 진영으로 나눌 때 슬픔의 총합이 최소가 되게 하는 문제이다. 각 사람들을 하나의 정점으로, 사람 간 슬픔의 정도를 간선의 용량이라 생각했을 때, 이 문제는 정점 집합 \(A\)에서 집합 \(B\)로 연결된 용량의 최솟값을 찾는 문제로 바꿀 수 있다. 즉 최대 유량 최소 컷(MFMC) 문제로, 최대 유량을 찾고 집합을 두 개로 나누면 문제를 해결할 수 있다.

이때, 에드몬드-카프 알고리즘으로 해결하면 시간 초과가 발생하므로 디닉 알고리즘을 써야만 한다.

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef tuple<int, int, int> tiii;

const int MAXN = 502;
const int INF = 1e9;
int N;
int cap[MAXN][MAXN], fl[MAXN][MAXN];
vector<int> graph[MAXN];
int dist[MAXN];
int idx[MAXN];
bool chk[MAXN];
int S = 0, T = MAXN-1;

bool BFS()
{
    memset(dist, -1, sizeof(dist));
    queue<int> qu;
    dist[S] = 0; qu.push(S);
    
    while(!qu.empty())
    {
        int u = qu.front(); qu.pop();
        for(int v : graph[u])
        {
            if(dist[v] == -1 && cap[u][v] - fl[u][v] > 0) //유량을 더 흘릴 수 있어야 한다.
            {
                dist[v] = dist[u] + 1;
                qu.push(v);
            }
        }
    }
    return dist[T] != -1;
}

int DFS(int u, int FL) //u : 정점, FL : 찾고 있는 증가 간선에서 흘릴 수 있는 유량
{
    if(u == T) return FL;
    //DFS를 여러 번 돌리기 때문에 간선을 중복해서 살펴보게 된다. 
    //이를 막기 위해 지금까지 살펴본 index에서부터 조사를 시작한다.
    //idx[u] : 지금까지 정점 u에서 살펴본 graph[u]의 index
    for(int& i=idx[u]; i<graph[u].size(); i++) 
    {
        int v = graph[u][i];
        if(dist[v] == dist[u] + 1 && cap[u][v] - fl[u][v] > 0)
        {
            int f = DFS(v, min(FL, cap[u][v] - fl[u][v]));
            if(f > 0)
            {
                fl[u][v] += f;
                fl[v][u] -= f;
                return f;
            }
        }
    }
    return 0;
}

void decom(int u)
{
    chk[u] = true; //chk[u] : source가 속한 집합인지 아닌지 확인
    for(int v : graph[u])
    {
        if(chk[v] || cap[u][v] == fl[u][v]) continue; 
        //u에서 v로 흘릴 수 있는 잔여 용량이 있다면 정점 v도 source가 속한 집합에 들어간다.
        decom(v);
    }
}

int main()
{
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	
    //그래프 구현
	cin >> N;
	for(int i=1; i<=N; i++)
	{
	    int x; cin >> x;
	    if(x == 1) 
	    {
	        cap[S][i] = INF;
	        graph[S].push_back(i);
	        graph[i].push_back(S);
	    }
	    else if(x == 2) 
	    {
	        cap[i][T] = INF;
	        graph[i].push_back(T);
	        graph[T].push_back(i);
	    }
	}
	for(int i=1; i<=N; i++)
	{
	    for(int j=1; j<=N; j++)
	    {
	        cin >> cap[i][j];
	        if(i != j) graph[i].push_back(j);
	    }
	}
	
	int ans = 0;
	while(BFS()) //BFS를 돌려 각 정점에 대해 S부터의 거리를 구한다.
	{
	    memset(idx, 0, sizeof(idx));
	    while(true) //BFS를 돌린 상태에서 증가 간선을 최대한 찾고, 유량을 최대한 흘린다.
	    {
	        int FL = DFS(S, INF);
	        if(FL == 0) break;
	        ans += FL;
	    }
	}
	cout << ans << "\n";
	
	decom(S); //집합 나누기
    
    for(int u=1; u<=N; u++) 
        if(chk[u]) cout << u << " ";
    cout << "\n";
    for(int u=1; u<=N; u++) 
        if(!chk[u]) cout << u << " ";
    cout << "\n";
}