1931 - Painting a Grid With Three Different Colors

info

• 문제 보기: 1931 - Painting a Grid With Three Different Colors
• 소요 시간: 💥1시간 초과
• 풀이 언어: java
• 체감 난이도: 4️⃣
• 리뷰 횟수: ✅

풀이 키워드

스포주의

DP 비트마스킹 dfs 조합론

풀이 코드

info

• 메모리: 55170 KB
• 시간: 255 ms

import java.util.*;

class Solution {
    final int MOD = 1_000_000_007;
    int m;
    int n;
    List<int[]> mColList;
    List<List<Integer>> nxtColIdList; // list of mColList id list
    Map<Integer, Integer> memo;

    public long dp(int x, int prevColId) {
        if (x == n) return 1;

        int key = x << 10 | prevColId;
        if (memo.containsKey(key)) return memo.get(key);

        long res = 0L;
        for (int id : nxtColIdList.get(prevColId)) {
            res = (res + dp(x+1, id)) % MOD;
        }

        memo.put(key, (int)res);
        return res;
    }

    public int colorTheGrid(int m, int n) {
        this.m = m;
        this.n = n;
        mColList = new ArrayList<>();
        nxtColIdList = new ArrayList<>();
        memo = new HashMap<>();

        // 1. dfs: pre-compute m-length columns
        initMColList(new ArrayList<>(), 0);

        // 2. pre-compute valid adj next columns for each m-length column
        initNxtColIdList();

        // 3. calc # of ways
        long ans = 0L;
        for (int id = 0; id < mColList.size(); ++id) {
            ans = (ans + dp(1, id)) % MOD;
        }

        return (int) ans;
    }

    public void initMColList(List<Integer> data, int y) {
        if (y == m) { // base case
            int[] col = new int[m];
            for (int i = 0; i < m; ++i) col[i] = data.get(i);
            mColList.add(col);
            return;
        }

        for (int c = 0; c < 3; ++c) { // iterate RGB
            if (0 < y && data.get(y-1) == c) continue;
            data.add(c);
            initMColList(data, y+1);
            data.remove(data.size()-1);
        }
    }

    public void initNxtColIdList() {
        for (int aid = 0; aid < mColList.size(); ++aid) {
            nxtColIdList.add(aid, new ArrayList<>());
            for (int bid = 0; bid < mColList.size(); ++bid) {
                if (check(aid, bid)) {
                    nxtColIdList.get(aid).add(bid);
                }
            }
        }
    }

    // check if b can be next to a
    public boolean check(int aid, int bid) {
        int[] a = mColList.get(aid);
        int[] b = mColList.get(bid);

        for (int i = 0; i < m; ++i) {
            if (a[i] == b[i]) return false;
        }

        return true;
    }
}

풀이 해설

단순 백트래킹은 TLE 터지는데 DP로 접근하자니 메모이제이션 적용하기 빡센 유형이다.

왜냐하면 메모이제이션을 위해 DP state를 '적당하게' 정의해야 하는데

현재 칸, 또는 인접 칸(위쪽, 왼쪽)의 칠해진 색을 state 정의로 사용하면 WA이기 때문이다.

❌ 실패한 발상

WA1

memo[i][c] -> ❌

각 칸의 번호를 i로 정의할 때, 해당 칸에 색칠된 색을 c라고 하자.

특정 칸에 칠해진 색 뿐만 아니라 인접한 다른 칸의 색칠 상태가 경우의 수에 영향을 미치므로,

해당 state는 잘못된 정의이다.

WA2

memo[i][tc][lc] -> ❌

각 칸의 번호를 i로 정의할 때, 해당 칸의 위쪽에 칠해진 색을 tc, 왼쪽에 칠해진 색을 lc라고 하자.

이 역시 다른 칸에서 경우의 수에 영향을 주기 때문에 잘못된 state 정의이다.

😮 그럼 어쩌라고?

그러니까, 칸 단위가 아니라 한 줄 통째로 메모이제이션 해야 한다.

따라서 이 문제는 column(세로) 단위로 칠할 수 있는 모든 경우들을 pre-compute 해서

이 column들이 인접하게 놓일 수 있는 경우의 수를 x축 방향으로 dp를 수행하여 구하는 식이다.

물론 column끼리 인접하게 놓일 수 있는지도 pre-compute 해서 nxtColIdList로 기록해둔다.

(만약 n과 m 범위가 거꾸로 된다면, row를 pre-compute 해서 y축 방향으로 dp를 수행하면 된다.)

[DFS] pre-compute columns

public void initMColList(List<Integer> data, int y) {
    if (y == m) { // base case
        int[] col = new int[m];
        for (int i = 0; i < m; ++i) col[i] = data.get(i);
        mColList.add(col);
        return;
    }

    for (int c = 0; c < 3; ++c) { // iterate RGB
        if (0 < y && data.get(y-1) == c) continue;
        data.add(c);
        initMColList(data, y+1);
        data.remove(data.size()-1);
    }
}

pre-compute valid next-columns

public void initNxtColIdList() {
    for (int aid = 0; aid < mColList.size(); ++aid) {
        nxtColIdList.add(aid, new ArrayList<>());
        for (int bid = 0; bid < mColList.size(); ++bid) {
            if (check(aid, bid)) {
                nxtColIdList.get(aid).add(bid);
            }
        }
    }
}

// check if b can be next to a
public boolean check(int aid, int bid) {
    int[] a = mColList.get(aid);
    int[] b = mColList.get(bid);

    for (int i = 0; i < m; ++i) {
        if (a[i] == b[i]) return false;
    }

    return true;
}

범위랑 dp 방향이랑 무슨 상관?

pre-compute 때문에 m, n의 범위에 따라 메모이제이션 대상과 dp 방향이 달라진다.

가능한 길이 m의 모든 column을 구하려 할 때, 경우의 수는 $3^5 = 243$ 미만인데

만약 m의 범위가 n처럼 $\le 1000$ 이었다면 $3^{1000}$ 이 되니까... 처리하기 힘들어진다.

따라서 범위를 보고 대강 row 기준으로 할지 column 기준으로 할지 유추할 수 있다.

✨ State Compression DP

column의 색칠 상태를 단순히 0/1/2 조합의 m글자 String으로 처리해서

dp의 x좌표와 같이 묶어 해시로 메모이제이션 처리할 수도 있다. 👉 900ms 대 성능

하지만 mColList의 인덱스를 id로 잡아 이를 x좌표랑 같이 묶으면 200ms 대로 최적화가 가능하다.

$m \le 5$ 이므로 $id < 3^5 < 2^8$

$n \le 1000$ 이므로 $x \le 1000 < 2^{10}$ 이어서 같이 묶어도 int 범위 내로 비트마스킹 가능하다.

아니면 그냥 memo[x][prevColId]와 같이 2차원 메모이제이션을 하든가.. 근데 저번 문제의 PTSD가

public long dp(int x, int prevColId) {
    if (x == n) return 1;

    int key = x << 10 | prevColId;
    if (memo.containsKey(key)) return memo.get(key);
    ...

메모

• 카카오 공채 고난이도 정도면 나올수도
• 이 문제도 matrix exponentiation으로 접근 가능하다.
  • 하지만 실전적으로는 classic DP 방식이 우선시된다.