백준 18352 특정 거리의 도시 찾기 Rust

rust 연습 삼아 간단한 알고리즘 문제들을 하루 한 문제 정도 하기로 결정했습니다.

시간 날 때 rust로 사이드 프로젝트도 간단히 해볼 생각입니다. (아이디어 생각은 해놨지만 언제 할런지....)

rust로 ps 자체에 대해 접근하기는 꽤 까다롭습니다. 입출력 속도 문제, 타입 문제 등이 있어 몇 개 풀어본 결과 괜찮은 방식은

use std::io::{stdin, Read};
use std::fmt::Write;

fn main() {

    let mut input = String::new();  
    let mut output = String::new();  
    stdin().read\_to\_string(&mut input).unwrap();  

    let mut input = input.split\_ascii\_whitespace();

    // n,m 입력 받기  
    let n = input.next().unwrap().parse::<usize>().unwrap();  
    let m = input.next().unwrap().parse::<usize>().unwrap();

    // ~~ res 연산

    writeln!(output, "{}", res).unwrap();

    print!("{}", output);
}

이런 식으로 입출력을 한 번에 처리하는 게 가장 깔끔했습니다.

입력 시 next로 공백, 줄바꿈 기준으로 다음 글자 있는 지 판단하고, 출력 시, output에 기록된 모든 내용을 한 번에 출력시킵니다.

입력 시 필요한 type은 parse:: 으로 해당 type에 맞게 parsing 시켜주면 됩니다.

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

간단히 다익스트라 + 우선순위 큐 기반의 문제를 풀어보면 다음과 같습니다.

N이 최대 30만이므로, 노드 개수가 N일 때 다익스트라 + 우선순위 큐 이용 시 O(N log N) 이므로 충분히 통과합니다.

2d vector로 graph를 구성해도 되지만, 그냥 hashmap 으로 구성했습니다. 해시 계산이 없는 2d vector가 더 빠르지만 연습삼아 hashmap으로 접근했습니다.

다익스트라 + 우선순위 큐 로직은 다음과 같습니다.

1. 전체 dist 배열을 INF로 초기화

2. 그래프 구성 시 (연결된 노드 번호, 엣지 가중치) 로 구성

3. 우선순위 큐에 (0, 시작점) 넣기 (min heap으로 구성)

4. 시작점으로부터 탐색 시작

5. 큐에서 하나씩 빼면서

6. 연결된 노드들 중 방문 안 했고, 현재까지 체크된 길이와 현재 엣지가 업데이트될 수 있으면 업데이트시키고 큐에 넣고 방문 체크

7. 5,6번 반복

코드는 다음과 같습니다.

use std::io::{stdin, Read};
use std::fmt::Write;
use std::collections::{BinaryHeap, HashMap};
use std::cmp::Reverse;

// n : 노드 개수, start : 시작점, graph : 그래프 정보
// return : 거리 정보 
fn dijkstra(n: usize, start: usize, graph: HashMap<usize, Vec<(usize, usize)>>) -> Vec<usize> {
    // 거리 초기화
    let mut dist: Vec<usize> = Vec::new();
    let mut visited: Vec<bool> = Vec::new();
    for _ in 0..n {
        dist.push(usize::MAX);
        visited.push(false);
    }

    // 우선순위 큐 : Reverse로 min heap 으로 동작시킴
    // pq 원소 : (거리, 노드)
    let mut pq: BinaryHeap<Reverse<(usize, usize)>> = BinaryHeap::new();
    pq.push(Reverse((0usize, start)));
    visited[start] = true;

    // 우선순위 큐 + 다익스트라 
    while !pq.is_empty() {
        let ele = pq.pop().unwrap().0;
        let cost = ele.0;
        let node = ele.1;

        if let Some(v) = graph.get(&node) {
            for ele in v.iter() {
                let next_node = (*ele).0;
                let next_cost = cost + (*ele).1;

                if visited[next_node] {
                    continue;
                }

                if next_cost < dist[next_node] {
                    dist[next_node] = next_cost;
                    pq.push(Reverse((next_cost, next_node)));
                    visited[next_node] = true;
                }
            }
        }
    }

    dist
}

fn main() {
    let mut input = String::new();
    let mut output = String::new();
    stdin().read_to_string(&mut input).unwrap();

    let mut input = input.split_ascii_whitespace();

    let n = input.next().unwrap().parse::<usize>().unwrap();
    let m = input.next().unwrap().parse::<usize>().unwrap();
    let k = input.next().unwrap().parse::<usize>().unwrap();
    let x = input.next().unwrap().parse::<usize>().unwrap() - 1;

    // 문제 상 거리 항상 1
    let mut graph: HashMap<usize, Vec<(usize, usize)>> = HashMap::new();
    for _ in 0..m {
        let a = input.next().unwrap().parse::<usize>().unwrap() - 1;
        let b = input.next().unwrap().parse::<usize>().unwrap() - 1;

        if !graph.contains_key(&a) {
            graph.insert(a, vec![]);
        } 
        if let Some(v) = graph.get_mut(&a) {
            v.push((b, 1usize));
        }
    }

    // 계산 
    let dist = dijkstra(n, x, graph);
    let mut res: Vec<usize> = Vec::new();

    for i in 0..dist.len() {
        if dist[i] == k {
            res.push(i+1);
        }
    }

    if res.len() == 0 {
        writeln!(output, "{}", -1).unwrap();
    } else {
        for r in res.iter() {
            writeln!(output, "{}", *r).unwrap();
        }
    }

    print!("{}", output);
}