Rust 빠른 입출력

수정

Rust로 문제 해결을 할 때 입출력은 항상 고민이 됩니다. std::io::stdin()을 매번 호출하면 락으로 인한 속도 저하가 발생하고, StdinLock::read_line을 매번 호출하기에는 입력이 줄로 구분되지 않고 공백으로 구분될 경우에 처리가 불편합니다. 여기서는 두 가지 텡플릿을 소개합니다. 예시 코드는 BOJ 15552 빠른 A+B입니다.

1. 짧은 템플릿

fn main() {
    use std::fmt::Write;
    let stdin = std::io::read_to_string(std::io::stdin()).unwrap();
    let mut stdout = String::new();
    let mut tokens = stdin.split_whitespace();
    let mut next = || tokens.next().unwrap();
    let t = next().parse().unwrap();
    for _ in 0..t {
        let a: i32 = next().parse().unwrap();
        let b: i32 = next().parse().unwrap();
        writeln!(stdout, "{}", a + b).unwrap();
    }
    print!("{stdout}");
}

tokens는 받은 입력을 공백 및 개행 기준으로 나눈 반복자입니다. (split_whitespace) 그리고 편의 클로저 next를 정의하는데, 이는 공백 및 개행으로 구분된 다음 문자열을 가져오는 역할을 합니다. 문제 해결 분야에서 보통 값이 공백 혹은 개행으로 구분되어 주어지기 때문에 이런 형식을 채택했습니다.

출력은 stdout에 모으고, 종료 시 print!를 해 줍니다.

2. println!을 그대로 사용할 수 있는 템플릿

fn solve(stdin: &str) {
    let mut tokens = stdin.split_whitespace();
    let mut next = || tokens.next().unwrap();
    let t = next().parse().unwrap();
    for _ in 0..t {
        let a: i32 = next().parse().unwrap();
        let b: i32 = next().parse().unwrap();
        println!("{}", a + b);
    }
}
 
fn main() {
    use std::io::*;
    solve(&read_to_string(stdin()).unwrap());
    STDOUT.with(|refcell| std::io::Write::flush(&mut *refcell.borrow_mut()).unwrap());
}
 
thread_local! {
    static STDOUT: std::cell::RefCell<std::io::BufWriter<std::io::StdoutLock<'static>>> = std::cell::RefCell::new(std::io::BufWriter::with_capacity(1 << 17, std::io::stdout().lock()));
}
 
#[macro_export]
macro_rules! println {
    ($($t:tt)*) => {
        STDOUT.with(|refcell| {
            use std::io::*;
            writeln!(refcell.borrow_mut(), $($t)*).unwrap();
        });
    };
}
#[macro_export]
macro_rules! print {
    ($($t:tt)*) => {
        STDOUT.with(|refcell| {
            use std::io::*;
            write!(refcell.borrow_mut(), $($t)*).unwrap();
        });
    };
}

우선 main을 살펴보면, 입력 전체를 받아 solve 함수에 넘겨주고 있습니다. solve 부분은 1번과 같지만 println!을 사용하여 입출력이 가능합니다. solve 함수가 종료되면 STDOUT을 flush해 줍니다.

아래의 매크로 부분은 println 및 print를 고속화하기 위한 장치입니다. Thread Local 변수인 STDOUT에 표준 출력을 연결한 BufWriter를 배정합니다. BufWriter의 용량인 1 << 17은 경험적으로 얻은 최적 크기입니다.

다음으로 println과 print가 STDOUT에 출력하도록 매크로를 재정의해줍니다. LocalKey<RefCell>에 with_borrow_mut이 존재하지만 많은 환경에서 지원하지 않아 풀어서 썼습니다.

Rust로 문제 해결 시에 다양한 입출력 인터페이스를 구상하고 만들어 봤지만, 위 두 인터페이스가 가장 간단하고 쓰기 편한 것 같습니다. 이 템플릿이 아무쪼록 편안한 문제 해결에 도움이 되길 바랍니다. 읽어주셔서 감사합니다.