백준 20938번 반짝반짝

문제 요약

전구 N개가 일자로 놓여있다. 만약 i번째 전구가 고장나면 i번째 전구를 포함하고 i번째 전구와 연결된 오른쪽 전구도 모두 꺼진다.

각 전구가 꺼질 확률과 구간 수 K가 주어진다.

N개의 전구를 K개의 연속 구간으로 쪼개려고 한다. 이 때, 켜진 전구의 개수의 기댓값의 최댓값을 계산해야 한다.

N은 최대 2500이며, K는 최대 10이다.

풀이

먼저, $C(i,j)$를 [i,j]에서 켜진 전구들의 개수의 기댓값이라고 하자. 그러면 원하는 값은 다음과 같은 DP로 구할 수 있다.
$$
\begin{gather}
dp(i, j) = \text{첫번째 전구부터 i번째 전구까지를 j개의 구간으로 쪼갰을 때 최댓값} \\
dp(i, j) = \underset{k<i}{max}(dp(k, j-1) + C(k+1, i))
\end{gather}
$$
$C(i,j)$가 전처리 되어있다면 위 dp 테이블은 $O(N^2K)$에 전부 채울 수 있고 우리가 원하는 값은 dp[N][K]가 된다.

이제 모든 $C(i,j)$를 적절한 시간 내에 계산해야 한다. 일단 계산해야할 구간은 $O(N^2)$가 된다.

먼저, [1,N]의 기댓값을 계산한다고 생각하자. 그러면 $N$개가 켜질 확률, $N-1$개가 켜질 확률, ... , 1개가 켜질 확률을 다 구하고 켜진 전구의 수와 곱해야 하는데 이렇게 접근하자니 식을 뽑기가 좀 힘들거 같다.

i번째 전구가 켜질확률을 생각해보자. 이 전구가 켜질 확률을 $p_i'$라고 하자.

$i$번째 전구가 켜졌다면 켜진 전구의 수는 1이고 그 확률이 $p_i'$기 때문에 $p_i'$의 합이 해당 구간의 기댓값이 될 것이다.

$i$번째 전구가 켜지기 위해선 본인과 연결된 전구 중에서 모든 왼쪽 전구가 고장나면 안된다.

그럴 확률은 $\prod_{k=1}^i (1-p_k)$인데 말로 풀어설명하면 나랑 연결된 모든 왼쪽 전구가 고장안나고 나도 고장안날 확률이다.

따라서, [1, N]에서 켜진 전구의 개수의 기댓값은 아래와 같이 표현할 수 있다.
$$
\begin{gather}
C(1,N) = \sum_{i=1}^{N}{\prod_{j=1}^{i}(1-p_j)} \\
C(s,e) = \sum_{i=s}^{e}{\prod_{j=s}^{i}(1-p_j)}
\end{gather}
$$
따라서, 모든 $C(i,j)$값을 $O(N^2)$에 구하는 것이 가능하고, 이를 통해서 처음에 만든 dp식을 이용하면 문제를 해결할 수 있다.

#include<bits/stdc++.h>
using namespace std;
using ld = long double;

ld C[2505][2505];
ld psum[2505][2505];
int N, K;
ld a[2505];
ld dp[2505][15];

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> K;
    for(int i=1;i<=N;++i) cin >> a[i];
    for(int i=1;i<=N;++i) {
        ld p = 1;
        ld sum = 0;
        for(int j=i;j<=N;++j) {
            p *= (1-a[j]);
            C[i][j] = p;
            sum += p;
            psum[i][j] = sum;
        }
    }
    for(int i=1;i<=N;++i) for(int j=1;j<=N;++j) dp[i][j] = 0;
    for(int i=1;i<=N;++i) dp[i][1] = psum[1][i];
    for(int j=2;j<=K;++j) {
        for(int i=j;i<=N;++i) {
            for(int k=j-1;k<i;++k) {
                dp[i][j] = max(dp[i][j], dp[k][j-1] + psum[k+1][i]);
            }
        }
    }
    cout << fixed << setprecision(15) << dp[N][K] << '\n';
    return 0;
}