[백준 18288] 팀 연습 - JAVA

문제 링크

문제 정의

문제 1번부터 N번까지 순서대로 각 문제를 A, B, C 중 한 명에게 배정할 때
다음 조건을 모두 만족하는 배정 방법의 수를 구하기

• $A$가 해결한 문제의 수는 $K$의 배수
• $B$는 연속해서 문제를 해결 불가
• $C$는 최소 한 문제 이상 해결

정답은 가능한 모든 배정 방법의 수를 $1,000,000,007$로 나눈 나머지

---

접근 방법

이 문제를 봤을 때 $O(N^2)$의 알고리즘으로는 풀 수 없다는 것을 파악했고
경우의 수를 구하는 문제이기 때문에 다이나믹 프로그래밍을 떠올렸습니다.

어떤 값을 기억해야할까?

그래서 어떤 값을 기억하고 있어야 하는지 생각헀습니다.

• $A$가 해결한 문제의 수는 $K$의 배수 → 문제를 푼 수를 $K$로 %한 값
• $B$는 연속 문제 해결 불가 → $B$가 현재 문제를 풀었는지에 대한 상태
• $C$는 최소 한 문제를 해결 → $C$가 한 문제를 풀었는지에 대한 상태

이 세 가지의 상태 값을 기억해야 하기 떄문에
저는 3차원 배열을 선언했습니다.

dp[i][j][k]의 의미는

• 직전에 누가 풀었는지($i$)
• $A$의 개수 $mod K(j)$
• $C$ 사용 여부($k$)를 상태로 갖는 경우의 수

---

1. 초기 상태 설정

첫 번째 문제를 누가 푸느냐로 시작 상태를 생성

if (K > 0) {
    dp[0][1 % K][0] = 1; // A
}
dp[1][0][0] = 1; // B
dp[2][0][1] = 1; // C

• $A$를 선택하면 → $A$ 개수 $+1$
• $B$를 선택하면 → 그대로
• $C$를 선택하면 → $C$ 사용 여부 = $1$

---

상태 전이 (핵심 로직)

다음 문제를 풀 사람을 선택하면서 상태를 갱신

for (int j = 0; j < 3; j++) {
    for (int k = 0; k < ck; k++) {
        for (int l = 0; l < 2; l++) {

현재 상태:

• $j:$ 이전 사람
• $k: A mod K$
• $l: C$ 사용 여부

---

A 선택

if (K > 0) {
    int nk = (k + 1) % K;
    toggle[0][nk][l] += dp[j][k][l];
}

• $A$ 문제 수 증가 → $mod K$ 갱신
• 이전 사람 = $A$로 변경

---

B 선택

if (j != 1) {
    toggle[1][k][l] += dp[j][k][l];
}

• 이전 사람이 $B$이면 선택 불가
• $A mod$, $C$ 여부는 그대로 유지

---

C 선택

toggle[2][k][1] += dp[j][k][l];

• $C$ 사용 여부를 무조건 $1$로 갱신

---

최종 정답 계산

for (int j = 0; j < 3; j++) {
    answer += dp[j][0][1];
}

조건:

• $A$의 개수 $mod K == 0 → k == 0$
• $C$를 최소 한 번 사용 → $l == 1$

---

코드

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

public class Main {
	static final int DIV = 1000000007;
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		int N = Integer.parseInt(st.nextToken());
		int K = Integer.parseInt(st.nextToken());

		System.out.println(dynamicProgramming(N, K));
	}

	private static int dynamicProgramming(int N, int K) {
		int ck = Math.max(K, 1);
		int[][][] dp = new int[3][ck][2];

		if (K > 0) {
			dp[0][1 % K][0] = 1;
		}
		dp[1][0][0] = 1;
		dp[2][0][1] = 1;

		for (int i = 0; i < N - 1; i++) {
			int[][][] toggle = new int[3][ck][2];
			for (int j = 0; j < 3; j++) {
				for (int k = 0; k < ck; k++) {
					for (int l = 0; l < 2; l++) {
						if (dp[j][k][l] == 0) continue;

						if (K > 0) {
							int nk = (k + 1) % K;
							toggle[0][nk][l] = (toggle[0][nk][l] + dp[j][k][l]) % DIV;
						}

						if (j != 1) {
							toggle[1][k][l] = (toggle[1][k][l] + dp[j][k][l]) % DIV;
						}

						toggle[2][k][1] = (toggle[2][k][1] + dp[j][k][l]) % DIV;
					}
				}
			}
			dp = toggle;
		}

		int answer = 0;
		for (int j = 0; j < 3; j++) {
			answer = (answer + dp[j][0][1]) % DIV;
		}

		return answer;
	}
}