[백준 - 5626번] 제단 - Java

문제 링크

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

문제

티어: 골드 1
알고리즘: dp

입력

첫째 줄에 제단의 열의 수 N이 주어진다. (1 ≤ N ≤ 10,000)

다음 줄에는 공백으로 구분된 N개의 정수 hi가 주어진다. (-1 ≤ hi ≤ 10,000) hi는 i번 열의 높이를 나타내며, -1인 경우에는 도둑이 그 열을 훔쳐간 상태를 나타낸다.

출력

첫째 줄에 가능한 제단의 수를 1,000,000,007로 나눈 나머지를 출력한다.

풀이

가능한 제단의 수를 1,000,000,007로 나눈 나머지를 출력해야 한다.

예제 입력 3번을 보면

6
-1 -1 -1 2 -1 -1

4번째를 2로 만들기 위해서는
3번째는 3, 2, 1이면서 5번째 또한 3, 2, 1이어야 한다.

그래서 3번째에 2가 있다고 가정해보면,
-1 -1 2 2 -1 -1이 되고, 3번째 2가 되려면 똑같이
2번째가 3, 2, 1이면서 4번째 또한 3, 2, 1이어야 한다.

그래서 4번째가 2일 때를 구하기 위해서는 3번째가 1, 2, 3일 때의 합이된다.
예를 들어 다음과 같다.
x x 1 2
x x 2 2
x x 3 2

여기서 5번째를 고려해주지 않아도 되는 이유는 결국 5번째를 구할 때도 4번째의 가능한 경우만을 더하기 때문에 문제될 게 없다. 예를 들어 5번째가 4라고 했을 때
x x x x 4
4 번째의 3, 4, 5만을 활용하게 된다.

다시 1 번째부터 보면
1 번째는 0만 들어갈 수 있다.
2 번째는 0또는 1이 들어갈 수 있다. -1이기 때문에 가능한 경우의 수를 전부 구해준다.
...
4 번째에서는 2가 들어가야 한다.
그래서 4번째가 2일 때 3번째는 1, 2, 3만 들어갈 수 있기 때문에 이 세 경우의 합을 넣어준다. 또한 4번째는 -1이 아닌 정해진 값이기 때문에 이 경우의 수만 넣어줘야 한다.
그래야 그 이후 2일 때의 경우의 수만으로 구하기 때문이다.
...
이후에는 N 번째까지 반복하고, N 번째가 0일 때의 경우의 수를 출력해주면 된다.
왜냐하면 N 번째는 0만 가능하기 때문이다.

이를 토대로 dp를 정의하면 2차원 배열이 필요하다. dp[A][B]
A는 몇 번째 인덱스인지
B는 그 제단의 높이를 의미한다. (B는 최대 5000을 가질 수 있다.)

그래서 dp[4][2]면 4번째가 2일 때 가능한 경우의 수를 뜻한다. (x x x 2)

이 풀이의 시간 복잡도는 O(N) * 5000이다.

소스 코드

import java.io.*;
import java.util.*;

public class Main {
    static final int MOD = 1000000007;
    static int N;
  public static void main(String args[]) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      N = Integer.parseInt(br.readLine());
      StringTokenizer st = new StringTokenizer(br.readLine());
      int[] arr = new int[N + 1];
      for(int i=1; i<=N; i++) {
          arr[i] = Integer.parseInt(st.nextToken());
          if(arr[i] > 5000) {
              System.out.println(0);
              return;
          }
      }
      if(arr[0] >= 1 || arr[N] >= 1) {
          System.out.println(0);
          return;
      }
      int[][] dp = new int[N + 1][5001]; //최대 5000이기 때문에
      dp[1][0] = 1;
      for(int i=2; i<=N; i++) {
          if(arr[i] == -1) {
              dp[i][0] = (dp[i - 1][0] + dp[i - 1][1]) % MOD;
              if(i == N) {
                  break;
              }
              int sum = (dp[i][0] + dp[i - 1][2]) % MOD;
              for(int j=1; j<=5000; j++) {
                  dp[i][j] = sum;
                  sum -= dp[i - 1][j - 1];
                  if(sum < 0) {
                      sum += MOD;
                  }
                  if(j + 2 <= 5000) {
                      sum = (sum + dp[i - 1][j + 2]) % MOD;
                  }
                  
              }
          } else {
              //-1이 아닌 경우는
              dp[i][arr[i]] = dp[i - 1][arr[i]];
              if(arr[i] - 1 >= 0) {
                  dp[i][arr[i]] = (dp[i][arr[i]] + dp[i - 1][arr[i] - 1]) % MOD;
              }
              if(arr[i] + 1 <=5000) {
                  dp[i][arr[i]] = (dp[i][arr[i]] + dp[i - 1][arr[i] + 1]) % MOD;
              }
          }
      }
      System.out.println(dp[N][0]);
  }
}