[백준] 4206번: 피보나치 단어 (D5) (Rust)

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

알고리즘

다이나믹 프로그래밍, 문자열, kmp

문제 요약

$F(n - 1)$ 과 $F(n - 2)$ 을 문자열로써 합치는 피보나치 단어 수열에 대해,
$F(n)$ 에 패턴 $p$ 가 몇 번 나오는지 구하는 문제입니다.

풀이

먼저, 머리를 쓰지 않는 풀이로 $F(0)$ 부터 $F(100)$ 까지의 수열을 전부 전처리하는 방법이 있겠습니다만..
메모리가 터져 나갑니다. $n$ 번째 피보나치 단어 수열의 길이는 $n$ 번째 피보나치 수와 정확히 같은데,
대충 근사해보면

$\text{len}(F(100)) \approx \cfrac{\varphi^{100}}{\sqrt 5}\approx354,224,848,179,261,915,075$

이런 말도 안 되는 길이가 나옵니다..
수열을 직접 만들 수는 없을 것 같습니다. 다른 방법을 생각해 봅시다.

$F(n)$ 은 $F(n-1)$ 과 $F(n-2)$ 를 합쳐서 만들어지는데,
$F(n-1)$ 에서 발견한 패턴 $p$ 의 개수와,
$F(n-2)$ 에서 발견한 패턴 $p$ 의 개수를 합해서
$F(n)$ 에서의 패턴의 개수를 알아낼 수 있지 않을까요?

$F(6)=1011010110110$ 을 예시로 살펴봅시다.
$F(5)=10110101$ , $F(4)=10110$ 입니다.

먼저 $10110101$ 에서 패턴을 찾고, $10110$ 에서 패턴을 찾아서 개수를 더합니다.
그런데, 이러면 계산하지 못하는 패턴이 있습니다.

예시로 패턴의 길이를 $4$ 라고 했을 때, 그림에서 저렇게 나뉘어지는 부분을 가로지르는 부분들은 검사하지 못합니다.
전체 수열은 필요 없고, 딱 저 가로지르는 부분만 필요합니다.
패턴의 길이를 $L$ 이라고 했을 때,

$F(n-1)$ 의 뒤에서 $L-1$ 개,
$F(n-2)$ 의 앞에서 $L-1$ 개만 추출해 합친 뒤,
해당 수열에서도 패턴을 찾아서 더하게 되면 $F(n)$ 에서의 패턴 개수를 구할 수 있습니다.
이를 위해선 각 수열마다 앞과 뒤의 정보를 계속 유지해주면 됩니다.

먼저 $F(0)$ 과 $F(1)$ 은 미리 저장해둡시다.

const F0: [bool; 1] = [false];
const F1: [bool; 1] = [true];

앞으로 쓸 일이 많으니, 편의상 $l=L-1$ 이라 하겠습니다.

$\text{dp}[i]=F(i)$ 에서의 패턴 개수
$\text{front}[i]=F(i)$ 의 앞의 최대 $l$ 개의 수
$\text{back}[i]=F(i)$ 의 뒤의 최대 $l$ 개의 수

이렇게 정의하고 시작하겠습니다.

let mut dp: Vec<u64> = vec![0; n + 1];
let mut front: Vec<Vec<bool>> = vec![Vec::new(); n + 1];
let mut back: Vec<Vec<bool>> = vec![Vec::new(); n + 1];
dp[0] = if p == F0 { 1 } else { 0 };
dp[1] = if p == F1 { 1 } else { 0 };
let l: usize = p.len() - 1;
if l > 0 {
    front[0] = F0.to_vec();
    front[1] = F1.to_vec();
    back[0] = F0.to_vec();
    back[1] = F1.to_vec();
}

초기화는 간단합니다. dp[0]과 dp[1]은 각각 p가 [false]인지 [true]인지만 검사하면 됩니다.
front[0~1], back[0~1]의 길이는 전부 $l$ 을 초과하면 안 됩니다.
$l$ 이 양수일 때만 각각 $F(0)$, $F(1)$ 로 초기화해줍니다.

for i in 2..=n {
    let ins: Vec<bool> = [back[i - 1].clone(), front[i - 2].clone()].concat();
    dp[i] = dp[i - 1] + dp[i - 2] + kmp_count(&ins, &p) as u64;
    
    let front_take: usize = l.min(front[i - 1].len());
    front[i] = [
        front[i - 1][..front_take].to_vec(),
        front[i - 2][..(l - front_take).min(front[i - 2].len())].to_vec(),
    ].concat();
    
    let back_take: usize = l.min(back[i - 2].len());
    back[i] = [
        back[i - 1][back[i - 1].len() - (l - back_take).min(back[i - 1].len())..].to_vec(),
        back[i - 2][back[i - 2].len() - back_take..].to_vec(),
    ].concat();
}

복잡해 보이지만 별거 없습니다.
$i$ 는 $2$ 부터 $n$ 까지 순서대로 증가합니다.

먼저 $F(i-1)$ 의 뒤에서 최대 $l$ 개의 수와, $F(i-2)$ 의 앞에서 최대 $l$ 개의 수를 합칩니다.
이렇게 만든 수열 ins가 위에서 말한 가로지르는 부분이 되겠네요.

$F(i-1)$ 에서의 패턴의 개수 dp[i - 1],
$F(i-2)$ 에서의 패턴의 개수 dp[i - 2],
ins에서의 패턴의 개수를 kmp로 구해 전부 더해주면
$F(i)$ 에서의 패턴의 개수 dp[i]가 완성됩니다.

이제 front[i], back[i]를 구할 차례입니다.

front[i]는 기본적으로 front[i - 1]에서부터 $l$ 개의 수를 가져와야 하는데,
front[i - 1]의 길이가 $l$ 보다 작을 경우, 남은 부분을 front[i - 2]에서 가져오면 됩니다.
남은 부분의 길이도 front[i - 2]를 초과할 경우, 그냥 거기까지만 가져오면 됩니다.
핵심은 front의 모든 배열의 길이를 최대 $l$ 로 제한하는 것이니까요.

back[i]도 방향만 정반대이고 거의 똑같습니다. 뒤에서부터 채우는데,
먼저 back[i - 2]의 뒤에서부터 가져오고,
남는 부분이 있으면 back[i - 1]의 뒤에서부터 또 가져온 다음 합치면 됩니다.
순서는 back[i -1] ‭→ back[i - 2]라는 것만 유의하시면 될 것 같습니다.

dp[n]을 출력하면 정답입니다.
n이 0 또는 1일 때는 따로 처리해주시면 됩니다.

코드(요약)

const F0: [bool; 1] = [false];
const F1: [bool; 1] = [true];

#[inline(always)]
fn solve(n: usize, p: Vec<bool>) -> u64 {
    if n == 0 { return if p == F0 { 1 } else { 0 }; }
    if n == 1 { return if p == F1 { 1 } else { 0 }; }
    let mut dp: Vec<u64> = vec![0; n + 1];
    let mut front: Vec<Vec<bool>> = vec![Vec::new(); n + 1];
    let mut back: Vec<Vec<bool>> = vec![Vec::new(); n + 1];
    dp[0] = if p == F0 { 1 } else { 0 };
    dp[1] = if p == F1 { 1 } else { 0 };
    let l: usize = p.len() - 1;
    if l > 0 {
        front[0] = F0.to_vec();
        front[1] = F1.to_vec();
        back[0] = F0.to_vec();
        back[1] = F1.to_vec();
    }
    for i in 2..=n {
        let ins: Vec<bool> = [back[i - 1].clone(), front[i - 2].clone()].concat();
        dp[i] = dp[i - 1] + dp[i - 2] + kmp_count(&ins, &p) as u64;
        let front_take: usize = l.min(front[i - 1].len());
        front[i] = [
            front[i - 1][..front_take].to_vec(),
            front[i - 2][..(l - front_take).min(front[i - 2].len())].to_vec(),
        ].concat();
        let back_take: usize = l.min(back[i - 2].len());
        back[i] = [
            back[i - 1][back[i - 1].len() - (l - back_take).min(back[i - 1].len())..].to_vec(),
            back[i - 2][back[i - 2].len() - back_take..].to_vec(),
        ].concat();
    }
    dp[n]
}

fn main() -> () {
    let mut i: u32 = 1;
    while let Some(n_str) = _r.next() {
        let n: usize = p!(n_str);
        let p: Vec<bool> = r!().as_bytes().iter().map(|&b| b == b'1').collect();
        wl!("Case {}: {}", i, solve(n, p));
        i += 1;
    }

    return;
}

전체 로직을 solve 함수로 분리해 작성하는 걸 상당히 싫어하는데, 문제의 특성상 어쩔 수 없이 이렇게 썼습니다..
설명을 보시면 알 수 있듯이, 사실 배열에 저렇게 전부 저장해놓을 필요가 없습니다.
이전의 정보와 그 이전의 정보만 유지하면 돼서, 단순히 변수를 두 개씩 사용하는 것으로도 대체할 수 있습니다.
그런데.. 메모리 제한이 넉넉해서 반도 안 쓰고 그냥 풀립니다. 더욱 깔끔한 풀이를 원하시는 분들은 최적화해서 작성하시면 되겠습니다.
수열을 합치는 과정에서도 복사를 남발하다 보니 속도도 매우 느립니다. (1위 2.4MB / 0ms)
많이 비효율적인 풀이라, 다른 분들의 풀이도 참고해 보시면 좋을 것 같습니다.

여담

오랜만에 되게 재미있는 문제였네요.
수열의 길이에 상관없이 공간 복잡도를 $O(l)$ 로 제한시키면서 구하는 과정이 상당히 인상적이었습니다.
사실 kmp 없이도 풀 수 있었을 것 같습니다.