LeetCode : 887. Super Egg Drop

문제

높이 $1$ 부터 $n$ 까지 총 $n$ 개의 계단이 있는 건물에서 $k$ 개의 동일한 계란을 떨어트리는 시행을 하려 한다.
이 때, 계란을 떨어트려도 깨지지 않는 최대 높이 $f$가 존재한다. ( $0 \le f \le n$ )

매 시행 온전한 계란을 임의의 높이 $x$ $(0 \le x \le n)$ 에서 떨어트리는데, 계란이 깨지지 않는다면 계란을 재활용할 수 있다.

$f$ 값을 확실하게 알아낼 수 있는 최소 시행횟수를 구하라.

제약 조건

• $1 \le k \le 100$
• $1 \le n \le 10^4$

예시

입력 : $k=1, n=2$
출력 : 2

• 계란을 높이 1 에서 떨어트려서 깨지면 $f$ 값은 0 이다.
• 계란이 깨지지 않는다면 높이 2 에서 떨어트려봐야 $f$ 값을 확정지을 수 있으므로 2 를 반환

풀이

우선, 계란에 대한 조건과 상관없이 $f$ 값을 효율적으로 찾기 위해서는 이분 탐색을 사용하는 것이 타당해보인다. 그런데 일단은 이분 탐색까지 생각하면 머리가 아프니 선형 탐색으로 구현을 시작했다.

그럼 이제 첫 시행에서 임의의 $x$ 를 잡고 계란을 떨어트린다고 가정.

1. 계란이 깨지는 경우

   $f \lt x$ 이므로, 깨지지 않은 나머지 계란 $k-1$ 개로 $x$ 보다 작은 수 $x-1$ 개를 탐색하면 된다.

2. 계란이 깨지지 않은 경우

   $f \gt x$ 이므로, 계란 $k$ 개로 $x$보다 큰 수 $n-x$ 개를 탐색하면 된다.

계란의 개수 $k$ 와 계단의 개수 $n$ 이 주어졌을 때 시행해야하는 횟수를 $dp[k][n]$ 이라고 한다면 두 개의 부분문제로 나누어서 표현할 수 있다. 이 때 두 경우의 최댓값을 최소화하는 $x$ 를 구해야 하므로 다음과 같이 표현할 수 있다.

$dp[k][n] = 1 + min_{x}(max(dp[k-1][x-1], dp[k][n-x]))$

1. 재귀

제일 먼저 기저 사례를 처리해야 하므로 어떤 것이 있을까 생각해봤다.

• 일단 계란이 하나밖에 없는 경우 ($k = 1$) 를 생각해본다면, $f$ 값을 확실하게 알아낼 수 있는 방법은 하나 하나 시도해보는 것이므로 시행 횟수는 $n$
• $n \lt 0$ 이면 $0$ , $n = 1$이면 $1$

그럼 이제 앞서 작성한 $dp$ 식과 함께 버무려서 코드를 작성하면,

class Solution {
public:
    int superEggDrop(int k, int n) {
        if(n < 1) return 0;
        if(k < 2) return n;
        
        // 최악의 경우는 n번 시행
        int min_num_trials = n;

		// 가능한 모든 x 값에 대해 
        for(int i=1; i<n; i++) {
            int num_trials = 1 + max(superEggDrop(k-1, i-1), superEggDrop(k, n-i));
            min_num_trials = min(min_num_trials, num_trials);
        }
        return min_num_trials;
    }
};

2. 메모이제이션

이미 계산했던 결과를 담을 배열 cache[n][k] 를 만들고,
다시 계산하지 않도록 메모이제이션 적용.

class Solution {
public:
    int cache[101][10001];
    Solution() {
        memset(cache, 0, sizeof(cache));
    }
    int superEggDrop(int k, int n) {
        if(n < 1) return 0;
        if(k < 2) return n;
       
       	// 이미 계산해봤는지 뒤져본다
        int& ret = cache[k][n];
        if(ret != 0) return ret;
        
        // 최악의 경우는 n번 시행
        int min_num_trials = n;

		// 가능한 모든 x 값에 대해 
        for(int i=1; i<n; i++) {
            int num_trials = 1 + max(superEggDrop(k-1, i-1), superEggDrop(k, n-i));
            min_num_trials = min(min_num_trials, num_trials);
        }
        // 답을 구했으면 cache에 저장하고 반환
        return ret = min_num_trials;
    }
};

그럼 이제 제출하면..
121개의 테스트케이스중에 69번 째, $k=5, n=10000$ 인 케이스에서 TLE가 발생한다.

3. 이분 탐색

메모이제이션을 적용하면 충분히 돌아갈 수 있을 것 같았는데 TLE가 발생했다.
불필요한 모든 $x$ 값까지 탐색하고 있기 때문이다.
따라서 이분 탐색을 적용해서 조금 더 최적화를 할 필요가 있다.
$\lfloor\cfrac{n}{2}\rfloor$ 인 곳에 공을 떨어트리고,

• 깨진 경우 : $1$ 부터 $\lfloor\cfrac{n}{2}\rfloor -1$ 까지 탐색
• 안깨진 경우 : $\lfloor\cfrac{n}{2}\rfloor +1$ 부터 n 까지 탐색

class Solution {
public:
    int cache[101][10001];
    Solution() {
        memset(cache, -1, sizeof(cache));
    }

    int superEggDrop(int k, int n) {
        if(k < 2) return n;
        if(n < 1) return 0;
        if(n == 1) return 1;

        // search cache
        int& ret = cache[k][n];
        if(ret != -1) return ret;

        int answer = n, low = 1, high = n, num_trials = 0;
        while(low <= high) {
            // determine range of x
            int middle = (low + high) / 2;
            int broken = superEggDrop(k-1, middle-1);
            int unbroken = superEggDrop(k, n-middle);
            num_trials = 1 + max(broken, unbroken);

            if(broken < unbroken) low = middle+1;
            else high = middle-1;

            answer = min(answer, num_trials);
        }

        // cache value and return
        return ret = answer;
    }
};

이제 제출하면 잘 돌아간다. 끝!