[백준] 24522번: 스네이크 게임 (P2) (Rust)

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

알고리즘

문자열, kmp

문제 요약

패턴 매칭 문제입니다. $3\le N,M\le500,000$ 이므로, $O(N+M)$ 정도에 풀어야 할 것 같습니다.
뱀 모양을 원본 문자열, 트리거 패턴을 찾으려는 문자열로 두고 KMP를 적용시키면 해당 시간 복잡도로 해결할 수 있습니다.
KMP 알고리즘에 대해서는 알고 있다는 전제 하에 설명하겠습니다.

풀이

좌표들로 주어지는 각 모양을 어떻게 변환해야 할지 생각해 봅시다.
좌표들의 절대적인 위치는 중요하지 않고, 얼마나 전진했는지, 어느 방향으로 회전을 했는지만 저장하면 됩니다.
좌표들을 순서대로 탐색하면서, [전진한 거리, 회전한 방향(왼쪽/오른쪽), 전진한 거리, ..., 전진한 거리] 형태로 저장하면 좋을 것 같습니다.
여기에서 유의할 점은, 전진한 거리는 1 이상의 정수이므로 회전한 방향을 나타낼 수는 0 이하의 정수 중에서 선택해야 합니다.
여기서는 왼쪽을 -1, 오른쪽은 0으로 설정하겠습니다.

방향을 결정하기 위해선, 먼저 두 좌표가 필요합니다.

저는 이렇게 나타내기로 했습니다. 두 좌표를 입력받고, 방향을 반환하는 함수를 만들어 봅시다.

const fn direction(y1: i32, x1: i32, y2: i32, x2: i32) -> u8 {
    if y1 > y2 {
        0
    } else if x2 > x1 {
        1
    } else if y2 > y1 {
        2
    } else { // if x1 > x2
        3
    }
}

대충 이렇게 되겠네요. (y, x) 형태의 좌표가 머리에 굳어버려서 이렇게밖에 생각이 안 납니다.. 죄송합니다.
그럼 이제 두 방향을 입력받고 회전 방향을 반환하는 함수를 생각해볼 수 있습니다.
왼쪽을 false, 오른쪽을 true로 표현하겠습니다.

const fn rotation(d1: u8, d2: u8) -> bool {
    match (d1, d2) {
        (0, 1) => true,
        (0, 3) => false,
        (1, 0) => false,
        (1, 2) => true,
        (2, 1) => false,
        (2, 3) => true,
        (3, 0) => true,
        (3, 2) => false,
        _ => unreachable!(),
    }
}

좀 지저분한데.. 실수하지 않기 위해서 그냥 전부 다 나열했습니다.
스네이크 게임의 특성상, 항상 왼쪽 또는 오른쪽으로만 회전하게 됩니다.
마지막으로 두 좌표를 입력받아서 이동 거리를 반환하는 함수를 만들면 됩니다.

const fn dist(y1: i32, x1: i32, y2: i32, x2: i32) -> i32 {
    (y2 - y1).abs() + (x2 - x1).abs()
}

항상 $y1=y2$ 이거나 $x1=x2$ 일 것이므로, 이런 식으로 간결하게 표현할 수 있습니다.

이제 입력된 좌표들을 해당 형태로 변환해 봅시다.
주어진 좌표의 개수가 $N$일 때, 선분의 개수는 $N-1$, 꼭짓점의 개수는 $N-2$ 이므로,
$2N-3$ 의 크기가 필요할 것이라고 예상할 수 있습니다.

let mut snake: Vec<i32> = Vec::with_capacity((n << 1) - 3);
let (sy, sx): (i32, i32) = (pr!(), pr!());
let (mut py, mut px): (i32, i32) = (pr!(), pr!());
let mut pd: u8 = direction(sy, sx, py, px);
snake.push(dist(sy, sx, py, px));
for _ in 0..n - 2 {
    let (y, x): (i32, i32) = (pr!(), pr!());
    let d: u8 = direction(py, px, y, x);
    snake.push(if rotation(pd, d) { 0 } else { -1 });
    snake.push(dist(py, px, y, x));
    pd = d;
    (py, px) = (y, x);
}

먼저 맨 앞의 두 좌표를 입력받습니다. 해당 두 좌표가 초기 방향을 결정하고, 초기 이동 거리를 배열에 넣으면서 시작하게 됩니다.
남은 $N-2$ 개의 좌표들을 입력받으면서, 이전 좌표를 통해 현재 방향을 계산하고, 이전 방향과 비교해 회전 방향을 삽입합니다.
이동 거리도 계산해서 넣어주고, 현재 방향과 좌표를 이전 방향과 좌표로써 저장해주면 반복을 돌면서 배열을 완성할 수 있습니다.

트리거 패턴도 완전히 같은 방식으로 완성하면 됩니다.

이제 패턴 매칭을 해 봅시다. 앞에서 언급한 대로, snake와 trigger 배열 모두
[전진한 거리, 회전한 방향(왼쪽/오른쪽), 전진한 거리, ..., 전진한 거리] 형태로 나타나게 됩니다.
그런데, 두 배열을 그대로 매칭하게 되면 답을 찾을 수 없습니다. 이유를 알아봅시다.

문제 본문에 좋은 예시가 있어서 그냥 이걸 사용하겠습니다.
예제 입력 2를 넣어서 snake와 trigger 배열을 디버깅해 보면,

$[8, 0, 4, 0, 2, 0, 2, -1, 2, -1, 4, -1, 8]$
$[1, -1, 1]$

이렇게 나오게 됩니다. 보자마자 알 수 있듯, 완전히 일치하는 패턴은 없습니다.
trigger 배열의 시작 거리와 끝 거리는 해당 위치에서의 snake 배열의 거리보다 작거나 같기만 하면 됩니다.
trigger의 시작과 끝을 제외한 배열을 패턴으로써 검색하고, 해당 조건을 만족하는지 검사하면 됩니다.
일치하는 인덱스들을 반환하는 함수 kmp_search에 대해,

let mut res: usize = 0;
let (l, r) = (trigger[0], *u!(trigger.last()));
for i in kmp_search(&snake, &trigger[1..trigger.len() - 1]) {
    if snake[i - 1] >= l && snake[i + trigger.len() - 2] >= r {
        res += 1;
    }
}

이렇게 해도, res는 여전히 3입니다.
현재 trigger 배열 $[1, -1, 1]$은 왼쪽으로 회전하는 경우입니다.
오른쪽으로 회전하는 경우, 즉 역순으로 처리한 배열도 검색해야 정확한 답이 나오게 됩니다.
trigger의 좌표들을 저장해서 역순으로 탐색해 배열을 완성하는 방법도 있으나,
단순히 정방향 배열을 거꾸로 보면서 방향만 반전시켜줘도 됩니다.

let trigger_rev: Vec<i32> = trigger.iter().rev().map(
    |&v| if v == 0 { -1 } else if v == -1 { 0 } else { v }
).collect();

현재 예시에서는 $[1, 0, 1]$이 되겠네요.

만약 trigger_rev가 trigger와 일치한다면 탐색하면 안 됩니다. 일치하지 않는 경우에만 계산해줍시다.

if trigger_rev != trigger {
    let (l, r) = (trigger_rev[0], *u!(trigger_rev.last()));
    for i in kmp_search(&snake, &trigger_rev[1..trigger_rev.len() - 1]) {
        if snake[i - 1] >= l && snake[i + trigger_rev.len() - 2] >= r {
            res += 1;
        }
    }
}

res를 출력하면 정답입니다.

코드 (요약)

#[inline(always)]
const fn direction(y1: i32, x1: i32, y2: i32, x2: i32) -> u8 {
    if y1 > y2 {
        0
    } else if x2 > x1 {
        1
    } else if y2 > y1 {
        2
    } else { // if x1 > x2
        3
    }
}

#[inline(always)]
const fn rotation(d1: u8, d2: u8) -> bool {
    match (d1, d2) {
        (0, 1) => true,
        (0, 3) => false,
        (1, 0) => false,
        (1, 2) => true,
        (2, 1) => false,
        (2, 3) => true,
        (3, 0) => true,
        (3, 2) => false,
        _ => unreachable!(),
    }
}

#[inline(always)]
const fn dist(y1: i32, x1: i32, y2: i32, x2: i32) -> i32 {
    (y2 - y1).abs() + (x2 - x1).abs()
}

fn main() -> () {
    let n: usize = pr!();
    let m: usize = pr!();

    let mut snake: Vec<i32> = Vec::with_capacity((n << 1) - 3);
    let (sy, sx): (i32, i32) = (pr!(), pr!());
    let (mut py, mut px): (i32, i32) = (pr!(), pr!());
    let mut pd: u8 = direction(sy, sx, py, px);
    snake.push(dist(sy, sx, py, px));
    for _ in 0..n - 2 {
        let (y, x): (i32, i32) = (pr!(), pr!());
        let d: u8 = direction(py, px, y, x);
        snake.push(if rotation(pd, d) { 0 } else { -1 });
        snake.push(dist(py, px, y, x));
        pd = d;
        (py, px) = (y, x);
    }

    let mut trigger: Vec<i32> = Vec::with_capacity((m << 1) - 3);
    let (sy, sx): (i32, i32) = (pr!(), pr!());
    (py, px) = (pr!(), pr!());
    pd = direction(sy, sx, py, px);
    trigger.push(dist(sy, sx, py, px));
    for _ in 0..m - 2 {
        let (y, x): (i32, i32) = (pr!(), pr!());
        let d: u8 = direction(py, px, y, x);
        trigger.push(if rotation(pd, d) { 0 } else { -1 });
        trigger.push(dist(py, px, y, x));
        pd = d;
        (py, px) = (y, x);
    }

    let trigger_rev: Vec<i32> = trigger.iter().rev().map(
        |&v| if v == 0 { -1 } else if v == -1 { 0 } else { v }
    ).collect();

    let mut res: usize = 0;

    let (l, r) = (trigger[0], *u!(trigger.last()));
    for i in kmp_search(&snake, &trigger[1..trigger.len() - 1]) {
        if snake[i - 1] >= l && snake[i + trigger.len() - 2] >= r {
            res += 1;
        }
    }
    
    if trigger_rev != trigger {
        let (l, r) = (trigger_rev[0], *u!(trigger_rev.last()));
        for i in kmp_search(&snake, &trigger_rev[1..trigger_rev.len() - 1]) {
            if snake[i - 1] >= l && snake[i + trigger_rev.len() - 2] >= r {
                res += 1;
            }
        }
    }

    wl!(res);

    return;
}

여담

블로그를 만들고 나서 제대로 올리는 첫 글입니다.
복붙을 목적으로 코드를 첨부하는 건 의미가 없다고 생각해서, 앞으로도 대략적인 코드만 주로 남겨놓을 것 같습니다.
읽어주셔서 감사합니다!