백준[16234] 인구 이동

문제 설명

문제 링크

http://icpc.me/16234

16234번: 인구 이동

풀이

dfs를 이용한 구현 문제다. bfs로 풀어도 똑같을 것 같다.

문제의 조건에 주어진대로 모든 점을 순회하면서 연합이 생기는지 확인하고 연합이 생긴다면 그 점들의 인구를 평균으로 바꿔주면 된다.

dfs를 순회하면서 연합이 되는 점들을 스택에 담아 두었다가 순회가 끝나고 스택에 있는 점들을 모두 평균으로 바꿔주는 방식으로 구현했다.

cf) 반복문에서 재귀를 호출할 때 retuen dfs(); 와같이 호출하여 반복문이 더 이상 돌지 않는 실수를 해서 시간을 많이 잡아먹었다. 이런 실수 은근 자주 하는 것 같다 주의하자.

코드

#include <cstdio>
#include <cmath>
#include <cstring>
#include <stack>

using namespace std;

int arr[55][55];
bool visited[55][55];
int dir[4][2] = {{1, 0}, {0, 1}, {-1, 0}, {0, -1}};
int n, l, r, sz, cnt, sum; 
stack<pair<int, int>> s;

int dfs(int y, int x) { 
    visited[y][x] = true;
    sz++;
    s.push({y, x});
    for (int i = 0; i < 4; i++) {
        int ny = y + dir[i][0], nx = x + dir[i][1];
        if (ny < 0 || ny >= n || nx < 0 || nx >= n || visited[ny][nx]) continue;
        if (abs(arr[y][x] - arr[ny][nx]) <= r && abs(arr[y][x] - arr[ny][nx]) >= l)  sum = arr[ny][nx] + dfs(ny, nx);
    }
    return sum;
}

int main() {
    scanf("%d %d %d", &n, &l, &r);
    for (int i = 0; i < n; i++) 
        for (int j = 0; j < n; j++) scanf("%d", &arr[i][j]);
    
    int ret = 0;
    do {
        sz = cnt = 0;
        memset(visited, 0, sizeof visited);
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                if (!visited[i][j]) {
                    sum = 0;
                    sum = dfs(i, j) + arr[i][j];
                    while (!s.empty()) {
                        auto [y, x] = s.top();
                        s.pop();
                        arr[y][x] = sum / sz;
                    }
                    if (sum != arr[i][j]) cnt++;
                    sz = 0;               
                }
            }
        }
        if (cnt) ret++;
    } while (cnt > 0);
    
    printf("%d", ret);
}