[BOJ 16234] 인구 이동 (Java, Python)

문제

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

16234번: 인구 이동

 

풀이

연합을 모두 찾은 후에 인구이동이 되어야 하므로, 찾은 연합을 저장해 둘 자료 구조를 생성한다.

1. 맵을 순회하면서 DFS로 연합을 찾아서 저장한다. 

2. 연합을 모두 찾은 후에, 연합 내 총 인구/나라 개수를 계산하여 각 나라의 인구를 갱신한다.

이 두 과정을 반복 한다. 연합을 찾았는데 연합의 개수가 0개인 경우 Loop를 종료한다. 

반복 횟수를 출력한다.

 

소스코드

package simulation;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class BOJ16234 {

    static class Union{
        int total;
        int average;
        ArrayList<Country> countries;

        Union(ArrayList<Country> countries){
            this.countries = countries;
            for(Country country: countries){
                total += map[country.x][country.y];
            }
            average = total /countries.size();
        }

        void updatePeopleCount(){
            for(Country country : countries){
                map[country.x][country.y] = average;
            }
        }
    }

    static class Country{
        int x;
        int y;

        Country(int x, int y){
            this.x = x;
            this.y = y;
        }
    }

    static int N, L, R;
    static int[][] map;
    static boolean[][] visited;
    static int[][] dir = new int[][]{{1, 0}, {-1, 0}, {0, 1}, {0, -1}};
    static ArrayList<Union> unions = new ArrayList<Union>();
    public static void main(String[] args) throws Exception  {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        L = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        map = new int[N][N];
        for(int i = 0; i < N; i++){
            st = new StringTokenizer(br.readLine());
            for(int j = 0; j < N; j++){
                map[i][j] = Integer.parseInt(st.nextToken());
            }
        }

        int day = 0;
        while(true){
            unions.clear();
            visited = new boolean[N][N];
            for(int i = 0; i < N; i++){
                for(int j = 0; j < N; j++){
                    if(visited[i][j]) continue;

                    visited[i][j] = true;
                    ArrayList<Country> countries = new ArrayList<Country>();
                    findUnion(i, j, countries);
                    if(countries.size() > 1){
                        unions.add(new Union(countries));
                    }
                }
            }
            for(Union union : unions){
                union.updatePeopleCount();
            }
            if(unions.size() == 0) break;
            day++;
        }

        System.out.println(day);

    }

    public static void findUnion(int x, int y, ArrayList<Country> countries){
        countries.add(new Country(x, y));
        for(int[] d : dir){
            int nx = x + d[0];
            int ny = y + d[1];
            if(isValid(map[x][y], nx, ny)){
                visited[nx][ny] = true;
                findUnion(nx, ny, countries);
            }
        }
    }

    public static boolean isValid(int val, int nx, int ny){
        if(nx < 0 || ny < 0 || nx >= N || ny >= N) return false;
        if(visited[nx][ny]) return false;
        int diff = Math.abs(val-map[nx][ny]);
        if(diff >= L && diff <= R) return true;
        return false;
    }

}

 

Python

Recursion Error와 80%에서 시간초과 

import sys

sys.setrecursionlimit(10**5) 

코드 추가 후 pypy3로 제출

 

import sys

sys.setrecursionlimit(10**5)

n, l, r = map(int, input().split())
board = []
for i in range(n):
    board.append(list(map(int, input().split())))

dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]
checked = [[False]*n for _ in range(n)]

def dfs(board, union, checked, x, y):
    checked[x][y] = True
    for i in range(4):
        nx = x + dx[i]
        ny = y + dy[i]
        if nx < 0 or ny < 0 or nx >= n or ny >= n or checked[nx][ny]:
            continue
        if l <= abs(board[x][y] - board[nx][ny]) <= r:
            union.append((nx, ny))
            dfs(board, union, checked, nx, ny)

day = 0
while True:
    unions = []
    checked = [[False] * n for _ in range(n)]

    for i in range(n):
        for j in range(n):
            if not checked[i][j]:
                union = [(i, j)]
                dfs(board, union, checked, i, j)
                unions.append(union)

    # board 인구수 업데이트
    if len(unions) == n*n:
        break
    else:
        day += 1
        for union in unions:
            people_count = 0
            for nation in union:
                x, y = nation
                people_count += board[x][y]
            updated_count = int(people_count//len(union))
            for nation in union:
                x, y = nation
                board[x][y] = updated_count

print(day)