프로그래머스 석유 시츄 [Lv. 2] (C++)

프로그래머스

N, M 크기의 격자 모양 땅 속에 석유가 발견 되었습니다. 시추관을 수직으로 단 하나만 뚫을 수 있을 때, 뽑을 수 있는 가장 많은 석유의 양을 구해주세요.

 

 

풀이 방법

시추관을 내릴 수 있는 모든 위치에서 바닥 까지 BFS를 사용하여 크기를 찾는 방법을 생각해 볼 수 있겠습니다. (세로 위치 0 에 시추관을 내린다면 0,0 ~ 0,N까지 모든 위치에서 BFS를 하여 석유 크기 추출) 하지만 이 방법은 시간 초과가 발생하기 때문에 안됩니다.

 

시간을 더 단축시킬 수 있는 방법은 Union - Find를 이용하는 것 입니다. 

우선 N * M의 속도로 BFS를 하여 land 배열에 저장 되어 있는 석유 덩러리를 1에서 석유 덩어리의 크기로 바꿔줍니다.

0 0 0 1 1 1 0 0 
0 0 0 0 1 1 0 0 
1 1 0 0 0 1 1 0 
1 1 1 0 0 0 0 0 
1 1 1 0 0 0 1 1 

🔽🔽🔽🔽🔽🔽

0 0 0 7 7 7 0 0 
0 0 0 0 7 7 0 0 
8 8 0 0 0 7 7 0 
8 8 8 0 0 0 0 0 
8 8 8 0 0 0 2 2

 

그와 동시에 각 석유 덩어리가 하나의 인덱스만을 바라보게끔 Union-Find를 해줍니다.

크기를 입력해주고 Union-Find로 하나의 인덱스를 바라보게끔 하는 이유는 매번 석유를 만날 때 마다 BFS로 크기를 체크하고 방문을 체크하는 과정을 생략해주기 위함입니다. (Union-Find가 바라보는 index가 set에 있다면 이미 체크한 석유 덩어리로 판단합니다)

 

이후 시추관을 내릴 수 있는 모든 위치에서 세로 방향으로 탐색을 시작하여, 0이 아닌 수를 찾으면 (석유가 있다면) 해당 인덱스의 부모 위치를 key로 하여 방문을 체크 -> 방문한적 없다면 크기를 기록, 방문 한적 있다면 건너뛰어가며 모든 위치에서 탐색을 해주면 됩니다.

 

 

정답 코드

 

#include <string>
#include <vector>
#include <queue>
#include <iostream>
#include <set>

using namespace std;

int Find(int b);
void UNION(int a, int b);

vector<int> parent;

int dx[4] = {0, 0, 1, -1};
int dy[4] = {1, -1, 0, 0};
void BFS(vector<vector<int>>& land, vector<vector<bool>>& visited, int i, int j)
{
    visited[i][j] = true;
    
    vector<pair<int, int>> v;
    queue<pair<int, int>> q;
    q.push({i, j});
    v.push_back({i, j});
    
    while(!q.empty())
    {
        auto[y, x] = q.front(); q.pop();

        UNION(i * land[0].size() + j, y * land[0].size() + x);
        
        for (int dir = 0; dir < 4; dir++)
        {
            int ny = y + dy[dir];
            int nx = x + dx[dir];

            // 시추 가능한 곳이 아니거나, 이미 시추한 곳이라면 continue;
            if (ny < 0 || ny >= land.size() || nx < 0 || nx >= land[i].size() 
                || visited[ny][nx] || land[ny][nx] == 0)
                continue;

            q.push({ny, nx});         
            v.push_back({ny, nx});
            visited[ny][nx] = true;
        }
    }
    
    for (auto& pos : v)
        land[pos.first][pos.second] = v.size();
}

int solution(vector<vector<int>> land) 
{
    parent.resize(land.size() * land[0].size());
    for (int i = 0; i < parent.size(); i++)
        parent[i] = i;
    

    
    vector<vector<bool>> visited(land.size(), vector<bool>(land[0].size()));  
    
    for (int i = 0; i < land.size(); i++)
        for (int j = 0; j < land[i].size(); j++)
            if (visited[i][j] == false && land[i][j] == 1)
                BFS(land, visited, i, j);
             

    int answer = 0;
    for (int j = 0; j < land[0].size(); j++)
    {
        int size = 0;
        set<int> visit;
        for (int i = 0; i < land.size(); i++)
        {
            int idx = Find(i * land[0].size() + j);
  
            if (land[i][j] == 0 || visit.find(idx) != visit.end())
                continue;
            visit.insert(idx);
            size += land[idx / land[0].size()][idx % land[0].size()];
        }
        
        answer = max(answer, size);
    }
    
    return answer;
}

void UNION(int a, int b)
{
    int A = Find(a);
    int B = Find(b);
    
    if (A <= B)
        parent[B] = A;
    else
        parent[A] = B;
}

int Find(int b)
{
    if (parent[b] == b)
        return b;
    
    return parent[b] = Find(parent[b]);
}