Rust로 PS하기 #1 - Fast I/O

작성일 | In

　PS를 하는 여러분들 중 대부분은, 이 문제를 한 번쯤은 풀어 봤을 것이라고 생각합니다. C++이나 python으로 짤 때는 별 문제 없이 맞았습니다!!를 받을 수 있었지만, Rust에서는 약간의 문제가 있었습니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
use std::io;

fn main(){
    let mut n = String::new();
    io::stdin()
        .read_line(&mut n)
        .expect("Fail");

    let n = n.trim().parse().expect("Fail");
    
    for i in 1..=n{
        println!("{}", i);
    }
}

　위 코드를 봅시다. 언뜻 보기엔 별 문제가 없어 보이는 코드지만, boj에 제출하면…

github

　시간 초과가 납니다. 이는 Rust의 I/O가 버퍼를 사용하지 않고, 매번 flush를 하기 때문에 발생한 문제입니다. 위 문제처럼 출력을 여러 번 하는 경우, println이 아닌, 버퍼를 사용하는 다른 출력 방법을 사용해야 합니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
use std::io;
use std::io::Write; // 1

fn main(){
    let stdout = io::stdout(); // 2
    let mut out = io::BufWriter::new(stdout.lock()); // 3

    let mut n = String::new();
    io::stdin()
        .read_line(&mut n)
        .expect("Fail");
    let n: i32 = n.trim().parse().expect("Fail");

    for i in 1..=n{
        writeln!(out, "{}", i); // 4
    }
}

　위 코드는 println 매크로로 인한 출력 속도 문제를 해결한 코드입니다. 처음 코드와의 차이점은 주석이 쳐져 있는 4개의 줄인데, 각각의 역할은 다음과 같습니다.

  1. 1번 주석은 writeln 매크로가 사용하는 함수인 std::io::Write의 함수 write_fmt를 사용할 수 있게 해줍니다.
  2. 2, 3번 주석은 버퍼 역할을 하는 BufWriter 타입의 변수 out을 선언합니다. 출력된 내용은 출력 버퍼인 out에 작성됩니다.
  3. 4번 주석에서 보이는 것처럼, println 매크로 대신 writeln 매크로를 사용합니다. 이때 writeln은 println과 달리, 첫 번째 인자로 출력 버퍼를 넘겨준다는 것을 알 수 있습니다.

github

　위 코드를 제출하면 12ms정도로 널널하게 맞았습니다!!를 받을 수 있습니다. 버퍼의 유무가 퍼포먼스에 얼마나 큰 영향을 미치는지 알 수 있는 문제였네요 :)

바보짓 해서 컴파일 에러 두 번 띄운건 비밀입니다.