BOJ#1626 두 번째로 작은 스패닝 트리

다야 캐야지 ~ 다이야 문제들의 특징은 풀이가 너무 이해하기 어렵다는 점인데, 이 문제는 선수문제(그래프와 MST)가 비교적 쉬운데, 저걸 풀면 바로 이해가 가며, 좀만 더 고치면 맞을 수 있다.

 

까다로운 반례 좀 있고, 추가된 테스트케이스가 많은 문제이다. 인터넷에 돌아다니는 블로그 중, 오래된 블로그 같은 경우에는 재채점 이후 틀린 코드가 꽤 있으니 주의하자.

(내 계정은 nemo이며, 이 글을 작성하는 시점에서는 맞았다.)

---

이 글에서는 그래프와 MST를 풀었다는 가정 아래 풀이를 설명하고자 한다. 안 풀었다면 풀러 가자. 어차피 짜야하는 부분이 같다. 당연하겠지만 가장 먼저 하게 되는 생각은 그래프와 MST의 정답을 정렬한 후 두 번째로 큰 답을 출력하면 되는 거 아닌가이다.

 

근데 몇 가지 조금 더 고려할 부분이 있다. 가중치가 0인 간선이 들어올 수 있다는 점이다. 이 경우 그래프와 MST의 코드는 가중치가 0이 아닌 간선을 잡아내지 못할 수 있다.

 

세 부분을 더 고려하자.

1. 우선 LCA에서 return의 초기값을 -1로 두자.

LCA에서 자기자신을 포함하는 간선과 가중치가 같지 않은 간선이 존재하지 않는 경우 -1을 반환하게 될 것이다.

2. LCA의 return을 갱신할 때 단순히 par[a][i] 해야하는 게 아니며, a부터 par[a][i]까지의 간선 중 자기자신을 포함하는 간선과 가중치가 같지 않은 간선으로 갱신해야 한다.

 

하나씩 짜보자.

1번(LCA부분)은 그리 복잡하지 않다. return만 -1로 고치며, ret를 갱신할 때 f(a, i, c)를 호출하여 2를 수행하게 할 것이다.

pair<int, ll> query(int a, int b, ll c) {
    ll ret = -1;
    if (dep[a] < dep[b]) swap(a, b);
    for (int i = tn - 1; i >= 0; i--)
        if (dep[a] - dep[b] >= (1 << i)) {
            ret = max(ret, f(a, i, c)), a = par[a][i].first;
        }
    if (a == b) return {a, ret};
    for (int i = tn - 1; i >= 0; i--)
        if (par[a][i].first != par[b][i].first)
            ret = max(ret, max(f(a, i, c), f(b, i, c))),
            a = par[a][i].first, b = par[b][i].first;
    return {par[a][0].first, max(ret, max(f(a, 0, c), f(b, 0, c)))};
}

 

f(a, i, c)는 a부터 a의 1<<i번째 부모까지의 간선 중 cost가 c가 아니면서 가장 큰 값을 반환하는 함수이다.

세그트리 짜듯이(?) 이분탐색으로 찾으면 된다.

 

ll f(int a, int b, ll c) {
    if (b == 0) return par[a][b].second == c ? -1 : par[a][b].second;
    ll ret = -1;
    if(par[a][b-1].second != c) ret = par[a][b-1].second;
    if(par[par[a][b-1].first][b-1].second != c) ret = max(ret, par[par[a][b-1].first][b-1].second);
    if(ret == 0) ret = max(f(a, b-1, c), f(par[a][b-1].first, b-1, c));
    return ret;
}

 

마지막으로 [그래프와 MST] 문제 부분에서, 현재 간선이 MST에 속하지 않는 경우 갱신하는 부분을, 현재 간선의 양 끝점의 경로에 cost가 다른 간선이 있는 경우만 정답 후보로 넣자.

 

for(int i=0; i<M; i++) {
    int isMST=0;
    int a=get<0>(E[i]), b=get<1>(E[i]); ll c=get<2>(E[i]);
    for(pair<int, ll> j : mst.mst[a]) {
        if(j.first == b && j.second == c) isMST = 1;
    }
    if(!isMST) {
        ll lq = lca.query(a, b, c).second;
        if(lq != -1) ans.emplace_back(mst.sum + c - lq);
    }
}

 

정답은 저기서 ans 배열 정렬한 후 두 번째로 작은 값 출력하면 된다.