백준 1626 두 번째로 작은 스패닝 트리

LCA - 최소 공통 조상을 사용하는 문제이다.

두 번째로 작은 스패닝 트리를 구할 때, MST에 포함되어있지 않은 에지를 골라서 그 에지가 잇는 두 노드의 경로 중 최대값 에지를 빼는 과정이 필요하다.

두 노드의 최단 경로를 찾을 때 LCA 알고리즘을 사용한다.

여기에, "두 번째로 작은" 이라는 조건이 붙으므로, 추가된 에지 비용과 경로 중 최대값 에지의 비용이 같으면 안된다. 그 다음으로 큰 에지를 골라야 한다.

 

여기까지 생각해보면, LCA에서 사용한 parents[K][V] 배열은 그대로 있어야 하고, path[K][V]를 만들어서 V에서 2^K 조상까지의 경로를 벡터형식으로 저장을 하면 좋을 것 같다. 나중에 두 노드에서 LCA까지 경로를 전부 아니까 거기서 MST 비용보다 큰 노드를 찾기만 하면 되겠지!

 

라는 생각으로 짰는데 메모리 초과가 되었다.

 

무식하게 벡터로 저장하는 방식이 너무 메모리를 많이 차지하는것 같다.

생각해보면 가장 큰값이나, 그 다음값만 가지고 있으면 되니 std::pair를 써서 처리하는게 좋아보인다.

pair로 처리를 하고 보니, pair와 pair를 비교해 가장 큰 값과, 두 번째로 큰 값을 구해주는 함수도 따로 만들어야 했다.

 

코드는 다음과 같다.

#include <iostream>
#include <queue>
#include <vector>
#include <algorithm>
#include <climits>
#include <string.h>

#define MAX_H 16
#define MAX_V 50000
#define PAIR std::pair<int, int>

int UFparents[MAX_V];
int pow2parents[MAX_H + 1][MAX_V + 1];
std::pair<int, int> pathList[MAX_H + 1][MAX_V + 1];
int depth[MAX_V];

typedef struct edge{
    int start;
    int end;
    int cost;
    int index;
    edge(int s, int e, int c, int i) : start(s), end(e), cost(c), index(i) {}
} edge;

struct cmp {
    bool operator()(edge e1, edge e2)
    {
        return e1.cost > e2.cost;
    }
};

void initLCA(int iCurNode, std::vector<std::pair<int, int> > adj[])
{
    for(int i = 0; i < adj[iCurNode].size(); i++)
    {
        int iNextNode = adj[iCurNode][i].first;
        int iNextEdge = adj[iCurNode][i].second;
        if (depth[iNextNode] == -1)
        {
            pow2parents[0][iNextNode] = iCurNode;
            depth[iNextNode] = depth[iCurNode] + 1;
            pathList[0][iNextNode] = PAIR(iNextEdge, -1);
            initLCA(iNextNode, adj);
        }
    }
}

int findParent(int a) {
    if (UFparents[a] == a)
        return a;
    return UFparents[a] = findParent(UFparents[a]);
}

void unionElement(int a, int b) {
    int pa = findParent(a);
    int pb = findParent(b);

    UFparents[pa] = pb;
}

PAIR getMax2Path (PAIR p1, PAIR p2)
{
    std::vector <int> temp;
    temp.push_back(p1.first);
    temp.push_back(p1.second);
    temp.push_back(p2.first);
    temp.push_back(p2.second);
    std::sort(temp.begin(), temp.end(), std::greater<int>());
    temp.erase(std::unique(temp.begin(), temp.end()), temp.end());
    PAIR ret;
    if (temp.size() > 2)
        ret = PAIR(temp[0], temp[1]);
    else if (temp.size() == 2)
        ret = PAIR(temp[0], -1);
    else
        ret = PAIR(-1, -1);
    return ret;
}

int main()
{
    int N, E;
    std::cin >> N >> E;
    for (int i = 0; i < N; i++)
        UFparents[i] = i;
    std::vector<std::pair<int, int> > adj[N];
    std::priority_queue<edge, std::vector<edge>, cmp> pq;
    std::vector<edge> vEdgeList;
    std::vector<int> vLeftOverEdge;
    for (int i = 0 ; i < E; i++)
    {
        int iNode1, iNode2, iCost;
        std::cin >> iNode1 >> iNode2 >> iCost;
        iNode2--;
        iNode1--;
        pq.push(edge(iNode1, iNode2, iCost, i));
        vEdgeList.push_back(edge(iNode1, iNode2, iCost, i));
    }
    int iMSTcost = 0;
    int iedgeCount = 0;
    while(!pq.empty())
    {
        edge curEdge = pq.top();
        pq.pop();

        if (findParent(curEdge.start) == findParent(curEdge.end))
        {
            vLeftOverEdge.push_back(curEdge.index);
            continue;
        }
        else
        {
            unionElement(curEdge.start, curEdge.end);
            iMSTcost += curEdge.cost;
            iedgeCount++;
            adj[curEdge.start].push_back(std::pair<int, int> (curEdge.end, curEdge.cost));
            adj[curEdge.end].push_back(std::pair<int, int> (curEdge.start, curEdge.cost));
        }
    }
    if (iedgeCount != N - 1)
    {
        std::cout << -1;
        return 0;
    }

    std::fill(depth, depth + N, -1);
    memset(pow2parents, -1, sizeof(pow2parents));
    depth[0] = 0;
    initLCA(0, adj);


    for (int K = 0; K < MAX_H - 1; K++)
    {
        for (int V = 1; V < N; V++)
        {
            if (pow2parents[K][V] != -1)
            {
                pow2parents[K+1][V] = pow2parents[K][pow2parents[K][V]];
                pathList[K+1][V] = getMax2Path(pathList[K][V], pathList[K][pow2parents[K][V]]);
            }
        }
    }
    // for (int K = 0; K < MAX_H; K++)
    // {
    //     for (int j = 0; j < N; j++)
    //     {
    //         std::cout << j << " node to 2^" << K << "anc : ";
    //         for (int i = 0; i < pathList[K][j].size() ; i++)
    //             std::cout << pathList[K][j][i] << " ";
    //         std::cout << "\n";
    //     }
    //     std::cout << "\n";
    // }

    //leftover edge 마다 두 노드 경로 중 최대값을 빼기. MST 와 만약 같다면 다음으로 큰 에지 빼기.
    // MST가 없을때. 2번째로 큰  MST 없을때.
    int iSecondMST = INT_MAX;
    for (int i = 0; i < vLeftOverEdge.size(); i++)
    {
        int iv = vEdgeList[vLeftOverEdge[i]].start;
        int iu = vEdgeList[vLeftOverEdge[i]].end;
        int edgeCost = vEdgeList[vLeftOverEdge[i]].cost;
        std::vector<int> trackPathList;
        PAIR costPair = PAIR(-1, -1);

        int iDepthDiff;
        if (depth[iv] < depth[iu])
            std::swap(iv, iu);
        iDepthDiff = depth[iv] - depth[iu];

        for (int bit = 0; iDepthDiff; bit++)
        {
            if (iDepthDiff % 2)
            {
                costPair = getMax2Path(costPair, pathList[bit][iv]);
                iv = pow2parents[bit][iv];
            }
            iDepthDiff /= 2;
        }

        if (iu != iv)
        {
            for (int K = MAX_H - 1; K >= 0; K--)
            {
                if (pow2parents[K][iu] != -1 && pow2parents[K][iu] != pow2parents[K][iv])
                {
                    costPair = getMax2Path(costPair, pathList[K][iu]);
                    costPair = getMax2Path(costPair, pathList[K][iv]);
                    iu = pow2parents[K][iu];
                    iv = pow2parents[K][iv];
                }
            }
            costPair = getMax2Path(costPair, pathList[0][iu]);
            costPair = getMax2Path(costPair, pathList[0][iv]);
        }
        if (costPair.first == -1 && costPair.second == -1)
            continue;
        if (costPair.first != -1 && iMSTcost + edgeCost - costPair.first != iMSTcost)
            iSecondMST = std::min(iSecondMST, iMSTcost + edgeCost - costPair.first);
        if (costPair.second != -1 && iMSTcost + edgeCost - costPair.second != iMSTcost)
            iSecondMST = std::min(iSecondMST, iMSTcost + edgeCost - costPair.second);
    }
    if (iSecondMST != INT_MAX)
        std::cout << iSecondMST;
    else
        std::cout << -1;
    
}

그래프와 MST 와 거의 유사하게 돌아가는데, 경로상 노드의 최대값, 두번째로 큰 값을 가지고 돌아다녀야 한다는 점이 가장 큰 차이점인것 같다. 

그 외에는 MST를 찾고 MST에 포함되지 않은 에지를 순회하면서 값을 찾는다는 점이 있겠다.

LCA 러시 주는 이정도로 하고, 다음주 부터는 다른 주제를 슬 보려고 한다...

 

오늘도 평온한 하루가 되길. 슨민.