백준[4013] ATM

문제 설명

문제 링크

https://www.acmicpc.net/problem/4013

풀이

SCC, 유니온 파인드, 위상 정렬, dp를 이용하여 풀이하였다.

다 푼 후에 다른 사람의 코드를 보고 깨달았지만, 유니온 파인드는 굳이 사용하지 않아도 된다.

출발지점부터 순회하며 dp를 돌리면 될 것 같지만 사이클이 존재하기 때문에 그럴 수 없다. 

그렇기 때문에 사이클이 없는 방향 그래프로 바꿔줘야 한다. 

사이클을 제거하기 위해 SCC를 사용한다. 아래 코드는 코사라주 알고리즘을 이용하였다.

구해진 SCC를 이용하여 사이클을 하나의 큰 노드로 합친다.

나는 이때 유니온 파인드를 사용하였다. 하지만 이미 SCC들을 모두 구했기 때문에 이미 합쳐져 있는 상태여서 유니온 파인드를 사용하지 않아도 된다. scc vector의 같은 인덱스에 들어있는 노드들끼리 합쳐져 있는 상태라고 볼 수 있다.

합치는 과정에서 레스토랑이 존재하는 사이클이라면 사이클 내의 모든 노드들은 레스토랑으로 갈 수 있기 때문에 사이클 내의 모든 노드를 레스토랑으로 만든다.  나는 이 부분을 놓쳐서 맞왜틀을 박았다.

scc끼리 모두 합친 후 합쳐진 노드들을 이용하여 새로운 그래프를 만든다.

그 후 출발지점부터 시작하여 dp를 돌리면 되는데 이때 방문하는 순서는 위상 정렬을 이용한다.

dp[다음 노드] = max(dp[다음 노드], dp[현재 노드] + cost[다음 노드])로 정의할 수 있다. 

이 문제는 최근에 공부한 알고리즘들을 모두 사용하기도 했고, 구현도 빡쎄서 재밌었다.

코드

사용된 변수가 워낙 많아 코드가 난잡하다. 천천히 읽어보자.

#include <cstdio>
#include <vector>
#include <stack>
#include <cstring>
#include <queue>

using namespace std;

int par[500005], cost[500005], dp[500005], indeg[500005];
bool visited[500005], restaurant[500005], check[500005];
vector<int> adj[500005], rev[500005], newG[500005], sc;
vector<vector<int>> scc;
stack<int> s;
queue<int> q;

int find(int x) {
    if (x == par[x]) return x;
    return par[x] = find(par[x]);
}

void merge(int x, int y) {
    x = find(x); y = find(y);
    if (x == y) return;
    if (restaurant[y]) restaurant[x] = true;
    par[y] = x;
    cost[x] += cost[y];
}

void dfs(int cur) {
    visited[cur] = true;
    for (auto nxt : adj[cur]) {
        if (!visited[nxt]) dfs(nxt);
    }
    s.push(cur);
}

void rev_dfs(int cur) {
    visited[cur] = true;
    for (auto nxt : rev[cur]) {
        if (!visited[nxt]) rev_dfs(nxt);
    }
    sc.push_back(cur);
}

int main() {
    int n, m; scanf("%d %d", &n, &m);
    for (int i = 0; i < m; i++) {
        int a, b; scanf("%d %d", &a, &b);
        adj[a].push_back(b);
        rev[b].push_back(a);
    }

    for (int i = 1; i <= n; i++) {
        scanf("%d", cost + i);
        par[i] = i;
    }

    int S, P; scanf("%d %d", &S, &P);
    for (int i = 0; i < P; i++) {
        int idx; scanf("%d", &idx);
        restaurant[idx] = true;
    }
	
    //scc구하기
    for (int i = 1; i <= n; i++) {
        if (!visited[i]) dfs(i);
    }

    memset(visited, 0, sizeof visited);
    while (!s.empty()) {
        int cur = s.top();
        s.pop();
        if (!visited[cur]) {
            rev_dfs(cur);
            scc.push_back(sc);
            sc.clear();
        }
    }

    // scc들끼리 뭉쳐서 하나의 정점으로 만들기
    for (auto &sc : scc) {
        for (int i = 0; i < sc.size() - 1; i++) {
            // merge 함수가 항상 x를 parent로 만드니까 S 노드가 없어지지 않도록 보존
            if (sc[i + 1] == S) swap(sc[i], sc[i + 1]);
            merge(sc[i], sc[i + 1]);
        }
    }

    for (int i = 1; i <= n; i++) {
        for (auto nxt : adj[i]) {
            int to = find(nxt);
            int from = find(i);
            // to와 from이 같으면 같은 scc이기 때문에 그래프 추가 안함
            if (to != from) {
                newG[from].push_back(to);
                indeg[to]++;
                // 정점이 합쳐져서 없어지는 노드가 생기기 때문에 존재하는 노드 번호 저장
                check[from] = check[to] = true;
            }
        }
    }

    for (int i = 1; i <= n; i++) {
        if (check[i] && !indeg[i]) q.push(i);
    }

    dp[S] = cost[S];
    bool flag = false;
    while (!q.empty()) {
        int cur = q.front();
        q.pop();
        if (cur == S) flag = true;
        for (auto nxt : newG[cur]) {
            // 출발 지점을 만나는 순간부터 dp 시작
            if (flag) dp[nxt] = max(dp[nxt], dp[cur] + cost[nxt]);
            if (--indeg[nxt] == 0) q.push(nxt);
        }
    }

    int ans = 0;
    for (int i = 1; i <= n; i++) {
        if (restaurant[i]) ans = max(ans, dp[i]); 
    }
    printf("%d", ans);
}