BOJ 17075 - 유물 복원 링크
(2024.03.25 기준 P5)

문제

$N \times M$ 크기의 석판의 각 칸의 정보가 주어진다. $1$이면 사람 한 명이 그려진 칸, $0$이면 빈 칸, $-1$이면 부서진 칸이다.

주어진 석판에서 가능한 모든 부분 직사각형에 대해 내부의 사람 수를 구하고, 그 수들을 모두 더해서 얻는 값이 $K$의 배수여야 할 때, 조건에 맞게 부서진 칸을 복원한 석판 하나를 출력

알고리즘

조합론 + 냅색 DP

풀이

모든 경우의 수를 따져보면 TLE다.

석판의 각 칸은 총합에 기여하는 정도가 정해져 있다. 다른 말로, 각 칸마다 모든 부분 직사각형에 포함되는 횟수는 정해져 있다. 이를 먼저 구해야 한다. 이는 BOJ 1286 문제에 쓰이는 조합론 공식을 사용해야 한다. 요약만 하면 $(i,j)$ 칸은 $(i+1) \times (N-i) \times (j+1) \times (M-j)$만큼 총합에 기여를 한다. (0-based index)
이에 대한 자세한 설명은 이 블로그에서 보자.

위에서 구한 가중치를 $w$라고 하자. 이제 우리는 총합을 $T$라고 했을 때, $T \bmod K = 0$인 경우를 구해야 한다. 이는 다음과 같은 점화식으로 냅색 dp로 풀어내면 된다.
$dp(i, j, k) = dp(pi, pj, k) ∨ dp(pi, pj, (k-w_{i,j}) \bmod K)$
전자는 $(i,j)$에 사람이 없는 경우, 후자는 $(i,j)$에 사람이 있는 경우이다.

그런데 단순히 가능한지만 판단하는게 아닌, 석판을 복원해야 한다. 그러므로 위에서 냅색 dp를 진행할 때, 참조한 경우가 전자인지 후자인지 따로 저장을 해놓으면 된다.

자세한 풀이는 공식 풀이를 참고하자.

코드

• C++

#include <bits/stdc++.h>
using namespace std;

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    int N, M, K; cin >> N >> M >> K;
    int S[N][M]; for (int i = 0; i < N; i++) for (int j = 0; j < M; j++) cin >> S[i][j];

    // w(i, j): i번째 줄의 j번째 칸에 사람이 있으면 총합에 더해지는 값
    int w[N][M];
    for (int i = 0; i < N; i++) for (int j = 0; j < M; j++)
        w[i][j] = (i + 1) * (N - i) * (j + 1) * (M - j);

    // dp(i, j, k): i번째 줄의 j번째 칸의 사람 유무를 결정해서 총합 mod K = k로 만들 수 있는가?
    // val(i, j, k): i번째 줄의 j번째 칸에서 총합 mod K = k가 가능할 때, 사람이 있어야 하면 1, 아니면 0.
    int dp[N + 1][M + 1][K]; memset(dp, 0, sizeof(dp));
    int val[N + 1][M + 1][K]; memset(val, -1, sizeof(val));
    int pi = 0, pj = M; // 직전 위치
    dp[pi][pj][0] = 1; // 아직 아무것도 하지 않은 상태
    for (int i = 1; i <= N; i++) for (int j = 1; j <= M; j++){
        for (int k = 0; k < K; k++){
            if (S[i - 1][j - 1] == 0){ // 현재 칸에서 사람이 없어야 한다면 직전 칸의 총합 mod K가 가능한지 확인
                if (dp[pi][pj][k]){
                    dp[i][j][k] = 1;
                    val[i][j][k] = 0;
                }
            }
            else if (S[i - 1][j - 1] == 1){ // 현재 칸에서 사람이 있어야 한다면 직전 칸의 총합-w(i,j) mod K가 가능한지 확인
                if (dp[pi][pj][((k - w[i - 1][j - 1]) % K + K) % K]){
                    dp[i][j][k] = 1;
                    val[i][j][k] = 1;
                }
            }
            else{ // 사람 유무를 정해야 하는 칸이라면, 두 경우에 대해 각각 체크한다.
                if (dp[pi][pj][k]){
                    dp[i][j][k] = 1;
                    val[i][j][k] = 0;
                }
                else if (dp[pi][pj][((k - w[i - 1][j - 1]) % K + K) % K]){
                    dp[i][j][k] = 1;
                    val[i][j][k] = 1;
                }
            }
        }
        pi = i; pj = j; // 현재 위치는 직전 위치가 된다.
    }

    // N번째 M칸까지 봤을 때, 총합 mod K = 0이 불가능하다면, K의 배수가 되는 총합이 없다는 뜻이다.
    if (!dp[N][M][0]){
        cout << -1;
        return 0;
    }

    // 역추적
    int ans[N][M], k = 0;
    for (int i = N; i > 0; i--) for (int j = M; j > 0; j--){
        if (val[i][j][k]){
            ans[i - 1][j - 1] = 1;
            k = ((k - w[i - 1][j - 1]) % K + K) % K;
            
        }
        else ans[i - 1][j - 1] = 0;
    }

    cout << 1 << '\n';
    for (int i = 0; i < N; i++){
        for (int j = 0; j < M; j++) cout << ans[i][j] << ' ';
        cout << '\n';
    }
}

• Python

import sys; input = sys.stdin.readline

N, M, K = map(int, input().split())
S = [list(map(int, input().split())) for _ in range(N)]

# w(i, j): i번째 줄의 j번째 칸에 사람이 있으면 총합에 더해지는 값
w = [[0] * M for _ in range(N)]
for i in range(N):
    for j in range(M):
        w[i][j] = (i + 1) * (N - i) * (j + 1) * (M - j)

# dp(i, j, k): i번째 줄의 j번째 칸의 사람 유무를 결정해서 총합 mod K = k로 만들 수 있는가?
# val(i, j, k): i번째 줄의 j번째 칸에서 총합 mod K = k가 가능할 때, 사람이 있어야 하면 1, 아니면 0.
dp = [[[0] * K for _ in range(M + 1)] for _ in range(N + 1)]
val = [[[-1] * K for _ in range(M + 1)] for _ in range(N + 1)]
pi = 0; pj = M # 직전 위치
dp[pi][pj][0] = 1 # 아직 아무것도 하지 않은 상태
for i in range(1, N + 1):
    for j in range(1, M + 1):
        for k in range(K):
            if S[i - 1][j - 1] == 0: # 현재 칸에서 사람이 없어야 한다면 직전 칸의 총합 mod K가 가능한지 확인
                if dp[pi][pj][k]:
                    dp[i][j][k] = 1
                    val[i][j][k] = 0
            elif S[i - 1][j - 1] == 1: # 현재 칸에서 사람이 있어야 한다면 직전 칸의 총합-w(i,j) mod K가 가능한지 확인
                if dp[pi][pj][(k - w[i - 1][j - 1]) % K]:
                    dp[i][j][k] = 1
                    val[i][j][k] = 1
            else: # 사람 유무를 정해야 하는 칸이라면, 두 경우에 대해 각각 체크한다.
                if dp[pi][pj][k]:
                    dp[i][j][k] = 1
                    val[i][j][k] = 0
                elif dp[pi][pj][(k - w[i - 1][j - 1]) % K]:
                    dp[i][j][k] = 1
                    val[i][j][k] = 1
        pi = i; pj = j # 현재 위치는 직전 위치가 된다.

# N번째 M칸까지 봤을 때, 총합 mod K = 0이 불가능하다면, K의 배수가 되는 총합이 없다는 뜻이다.
if not dp[N][M][0]:
    print(-1)
    exit()

# 역추적
ans = [[0] * M for _ in range(N)]
k = 0
for i in range(N, 0, -1):
    for j in range(M, 0, -1):
        print(i,j,k,val[i][j][k])
        if val[i][j][k]:
            ans[i - 1][j - 1] = 1
            k = (k - w[i - 1][j - 1]) % K

print(1)
for i in range(N):
    print(*ans[i])