1248번: Guess

문제

• 10부터 10까지의 정수로 이루어진 N개의 수열 A가 있습니다. 이 수열의 i번째 수부터 j번째 수까지의 합의 부호(+, -, 0)가 주어질 때, 부호 조건들을 모두 만족하는 수열 A를 찾는 프로그램을 작성하는 문제입니다.
• 입력: 첫째 줄에 수열의 크기 N (1 ≤ N ≤ 10)이 주어집니다. 둘째 줄에는 N(N+1)/2 길이의 문자열 형태의 부호 배열이 주어집니다.
• 출력: 부호 조건을 만족하는 수열 A를 공백으로 구분하여 출력합니다. (조건을 만족하는 수열이 여러 개일 경우 아무거나 하나만 출력합니다.)

풀이: 백트래킹을 활용한 부분 합 유효성 검사

이 문제는 10개의 자리에 -10부터 10까지 21개의 숫자를 무작위로 채우면 21^10 번의 탐색이 필요하여 시간 초과가 발생합니다. 따라서 백트래킹(DFS)을 이용해 조건을 만족하지 않는 경우는 빠르게 가지치기(Pruning)를 해야 합니다.

• 2차원 배열 할당: 주어진 1차원 부호 문자열을 직관적으로 다루기 위해 넉넉하게 n x n 크기의 2차원 배열 a에 저장합니다. a[i][j]는 ans[i]부터 ans[j]까지의 합의 부호를 의미하게 됩니다.
• 1차 가지치기 (자신의 부호 활용):a[idx][idx]는 ans[idx] 숫자 하나만의 부호를 의미합니다. 따라서 무작정 -10부터 10까지 모두 탐색하지 않고, 부호가 '+'면 1~10, '-'면 -10~-1, '0'이면 0만 탐색하도록 반복문의 범위를 좁힙니다.
• 2차 가지치기 (부분 합 실시간 검사):idx 위치의 숫자를 임시로 채웠다면, 다음 숫자를 찾으러 재귀 호출(go(idx + 1))을 넘어가기 전에 check 함수로 유효성을 검사합니다. i를 idx부터 0까지 거꾸로 내려가면서 sum에 ans[i]를 누적해 더하면, 이는 곧 ans[i] + ... + ans[idx]의 합이 됩니다. 이때의 합의 부호가 a[i][idx]의 부호와 하나라도 다르다면 즉시 false를 반환하여 잘못된 경로를 쳐냅니다.
• 프로그램 조기 종료: 정답 수열을 하나라도 찾았다면 더 이상 탐색할 필요가 없으므로 System.exit(0)을 호출하여 프로그램을 완전히 종료합니다.

코드

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

public class Main {
    static int n;
    static char[][] a;
    static int[] ans;

    public static void main(String[] args) throws IOException {
       BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
       BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

       n = Integer.parseInt(br.readLine());

       // 인덱스 초과 오류 방지 및 편의를 위해 N x N 배열로 선언
       a = new char[n][n];
       String s = br.readLine();
       int idx = 0;
       for (int i = 0; i < n; i++) {
          for (int j = i; j < n; j++) {
             a[i][j] = s.charAt(idx++);
          }
       }

       ans = new int[n];
       go(0);
    }

    static void go(int idx) {
       // N개의 숫자를 모두 채웠다면 정답을 출력하고 즉시 종료 (Base Case)
       if (idx == n) {
          for (int i = 0; i < n; i++) {
             System.out.print(ans[i] + " ");
          }
          System.out.println();
          System.exit(0);
       }

       char c = a[idx][idx];

       // a[idx][idx] 부호에 따라 탐색 범위를 한정 (가지치기 1)
       if (c == '0') {
          ans[idx] = 0;
          if (check(idx)) {
             go(idx + 1);
          }
       } else if (c == '+') {
          for (int i = 1; i <= 10; i++) {
             ans[idx] = i;
             if (check(idx)) {
                go(idx + 1);
             }
          }
       } else {
          for (int i = -10; i <= -1; i++) {
             ans[idx] = i;
             if (check(idx)) {
                go(idx + 1);
             }
          }
       }
    }

    // 방금 채운 ans[idx]가 포함된 부분합들이 부호 조건과 일치하는지 검증 (가지치기 2)
    static boolean check(int idx) {
       int sum = 0;

       // idx부터 0까지 거꾸로 더하면서 부분합을 확인
       for (int i = idx; i >= 0; i--) {
          sum += ans[i];

          if (a[i][idx] == '+' && sum <= 0) {
             return false;
          }
          if (a[i][idx] == '-' && sum >= 0) {
             return false;
          }
          if (a[i][idx] == '0' && sum != 0) {
             return false;
          }
       }

       return true;
    }
}

알고리즘 분류

브루트포스 알고리즘 (Brute Force), 백트래킹 (Backtracking)