atcoder/lib/
io.rs

1//! 空白区切り入力とバッファ付き出力を提供します。
2//!
3//! [`crate::input!`] は標準入力をスレッドごとに保持するため、複数回呼び出しても読み取り位置を
4//! 引き継ぎます。出力には、破棄時に flush する [`Writer`] を利用できます。
5
6use std::cell::RefCell;
7use std::fmt::Display;
8use std::io::{self, BufRead, BufWriter, Write};
9
10/// `BufRead` から空白区切りのトークンを読み取る scanner です。
11///
12/// 一度に1行ずつ読み込み、ASCII 空白文字で区切られたトークンを読みます。
13/// 読み取り位置は scanner が保持するため、同じ scanner への呼び出し間で続きから読み取れます。
14/// トークン読み取りの時間計算量はトークン長を `L` バイトとして償却 `O(L)` です。
15/// 入力が尽きた場合、読み取りに失敗した場合、またはトークンを `T` に変換できない場合は panic します。
16#[doc(hidden)]
17pub struct Scanner<R: BufRead> {
18    reader: R,
19    buffer: String,
20    cursor: usize,
21}
22
23impl<R: BufRead> Scanner<R> {
24    /// 入力元 `reader` を指定して scanner を作ります。
25    ///
26    /// 読み取り元の現在位置から読み始めます。以後の入力取得は遅延して行います。
27    ///
28    /// # Examples
29    ///
30    /// ```
31    /// use atcoder::io::Scanner;
32    ///
33    /// let mut scanner = Scanner::new(&b"12 hello"[..]);
34    /// assert_eq!(scanner.token::<usize>(), 12);
35    /// assert_eq!(scanner.token::<String>(), "hello");
36    /// ```
37    pub fn new(reader: R) -> Self {
38        Self {
39            reader,
40            buffer: String::new(),
41            cursor: 0,
42        }
43    }
44
45    /// 入力元 `reader` を指定して scanner を作ります。
46    ///
47    /// [`new`](Self::new) と同じです。
48    pub fn with_reader(reader: R) -> Self {
49        Self::new(reader)
50    }
51
52    /// 次の空白区切りトークンを `T` として返します。
53    ///
54    /// `T` は [`std::str::FromStr`] を実装する必要があります。トークンがなければ追加の入力を読み、
55    /// 入力終端に達したとき、または変換に失敗したときは panic します。
56    pub fn token<T: std::str::FromStr>(&mut self) -> T {
57        loop {
58            let bytes = self.buffer.as_bytes();
59            while self.cursor < bytes.len() && bytes[self.cursor].is_ascii_whitespace() {
60                self.cursor += 1;
61            }
62            let start = self.cursor;
63            while self.cursor < bytes.len() && !bytes[self.cursor].is_ascii_whitespace() {
64                self.cursor += 1;
65            }
66            if start < self.cursor {
67                return self.buffer[start..self.cursor]
68                    .parse()
69                    .ok()
70                    .expect("failed to parse token");
71            }
72            self.buffer.clear();
73            self.cursor = 0;
74            let read = self
75                .reader
76                .read_line(&mut self.buffer)
77                .expect("failed to read input");
78            assert!(read > 0, "input exhausted");
79        }
80    }
81}
82
83thread_local! {
84    #[doc(hidden)]
85    pub static SCANNER: RefCell<Scanner<io::BufReader<io::Stdin>>> =
86        RefCell::new(Scanner::new(io::BufReader::new(io::stdin())));
87
88    #[doc(hidden)]
89    pub static WRITER: RefCell<BufWriter<io::Stdout>> =
90        RefCell::new(BufWriter::new(io::stdout()));
91}
92
93/// `input!` の内部で入力形式を展開します。
94///
95/// 通常の型に加え、タプル、固定長入力、`chars`、`usize1`、`isize1` を扱います。
96///
97/// 各 `name: type` の形式で変数を宣言しながら値を読みます。`mut` を付けた名前は可変変数になります。
98/// `chars` はトークンを Unicode scalar value の `Vec<char>` にし、`[T; n]` は `n` 個を `Vec<T>`
99/// として読みます。`usize1` と `isize1` は入力整数から 1 を引くため、入力値は 1 以上である必要があります。
100/// タプル型 `(T1, T2, ...)` は各要素を順に読み取ります。
101#[doc(hidden)]
102#[macro_export]
103macro_rules! read_value {
104    ($scanner:expr, ($($ty:tt),*)) => {
105        ($( $crate::read_value!($scanner, $ty) ),*)
106    };
107    ($scanner:expr, [$ty:tt; $len:expr]) => {
108        (0..$len)
109            .map(|_| $crate::read_value!($scanner, $ty))
110            .collect::<Vec<_>>()
111    };
112    ($scanner:expr, chars) => {
113        $scanner.token::<String>().chars().collect::<Vec<char>>()
114    };
115    ($scanner:expr, usize1) => {
116        $scanner.token::<usize>() - 1
117    };
118    ($scanner:expr, isize1) => {
119        $scanner.token::<isize>() - 1
120    };
121    ($scanner:expr, $ty:ty) => {
122        $scanner.token::<$ty>()
123    };
124}
125
126/// 標準入力から変数を宣言しながら値を読み取ります。
127///
128/// # Examples
129///
130/// ```no_run
131/// atcoder::input! {
132///     n: usize,
133///     values: [i64; n],
134///     mut position: usize1,
135///     word: chars,
136/// }
137/// # let _ = (values, position, word);
138/// ```
139#[macro_export]
140macro_rules! input {
141    ($(,)?) => {};
142    (mut $($var:ident),+ : $ty:tt $(, $($rest:tt)*)?) => {
143        $(
144            let mut $var = $crate::io::SCANNER.with(|scanner| {
145                let mut scanner = scanner.borrow_mut();
146                $crate::read_value!(&mut *scanner, $ty)
147            });
148        )+
149        $($crate::input!($($rest)*);)?
150    };
151    ($($var:ident),+ : $ty:tt $(, $($rest:tt)*)?) => {
152        $(
153            let $var = $crate::io::SCANNER.with(|scanner| {
154                let mut scanner = scanner.borrow_mut();
155                $crate::read_value!(&mut *scanner, $ty)
156            });
157        )+
158        $($crate::input!($($rest)*);)?
159    };
160}
161
162/// グローバルなバッファへ改行なしで出力します。
163///
164/// 出力後は [`wflush`] を呼び出してください。呼び出しごとに標準出力へ直接書き込まず、
165/// スレッドローカルなバッファへ蓄積します。
166#[macro_export]
167macro_rules! wprint {
168    ($($argument:tt)*) => {
169        $crate::io::WRITER.with(|writer| {
170            std::io::Write::write_fmt(
171                &mut *writer.borrow_mut(),
172                format_args!($($argument)*),
173            )
174            .unwrap()
175        })
176    };
177}
178
179/// グローバルなバッファへ改行付きで出力します。
180///
181/// フォーマット構文は [`format_args!`] と同じです。実際の書き出しはバッファが満杯になるか、
182/// [`wflush`] が呼ばれたときに行われます。
183#[macro_export]
184macro_rules! wprintln {
185    () => {
186        $crate::wprint!("\n")
187    };
188    ($($argument:tt)*) => {
189        $crate::wprint!("{}\n", format_args!($($argument)*))
190    };
191}
192
193/// [`wprint!`] と [`wprintln!`] のスレッドローカルな標準出力バッファを flush します。
194///
195/// 競技プログラムの終了前や、対話型問題で出力を相手へ届ける必要があるときに呼び出します。
196/// 書き込みに失敗した場合は panic します。
197pub fn wflush() {
198    WRITER.with(|writer| writer.borrow_mut().flush().unwrap());
199}
200
201/// `Write` 先へ効率よく出力するヘルパです。
202///
203/// 値は内部の [`BufWriter`] に蓄積され、[`flush`](Self::flush) または `Drop` 時に
204/// 書き出されます。
205pub struct Writer<W: Write> {
206    writer: BufWriter<W>,
207}
208
209impl<W: Write> Writer<W> {
210    /// `writer` を出力先として作ります。
211    ///
212    /// 書き込みは内部バッファに蓄積されます。明示的に [`flush`](Self::flush) するか、値が破棄されると
213    /// バッファが出力先へ書き出されます。
214    pub fn with_writer(writer: W) -> Self {
215        Self {
216            writer: BufWriter::new(writer),
217        }
218    }
219
220    /// `Display` 可能な `value` を改行なしで出力バッファへ追加します。
221    ///
222    /// 書き込みに失敗した場合は panic します。
223    pub fn print(&mut self, value: impl Display) {
224        write!(self.writer, "{value}").unwrap();
225    }
226
227    /// `Display` 可能な `value` と改行を出力バッファへ追加します。
228    ///
229    /// 書き込みに失敗した場合は panic します。
230    pub fn println(&mut self, value: impl Display) {
231        writeln!(self.writer, "{value}").unwrap();
232    }
233
234    /// `condition` が `true` なら `Yes`、`false` なら `No` を改行付きで出力します。
235    pub fn print_yes_no(&mut self, condition: bool) {
236        self.println(if condition { "Yes" } else { "No" });
237    }
238
239    /// `Yes` を改行付きで出力します。
240    pub fn print_yes(&mut self) {
241        self.print_yes_no(true);
242    }
243
244    /// `No` を改行付きで出力します。
245    pub fn print_no(&mut self) {
246        self.print_yes_no(false);
247    }
248
249    /// `values` の各要素を `separator` で連結し、末尾に改行を出力します。
250    ///
251    /// 空イテレータの場合も改行を1つ出力します。イテレータの要素数を `N`、値と区切り文字の
252    /// 合計出力バイト数を `S` とすると、時間計算量は `O(N + S)` です。
253    pub fn join<T: Display>(&mut self, values: impl IntoIterator<Item = T>, separator: &str) {
254        let mut values = values.into_iter();
255        if let Some(first) = values.next() {
256            self.print(first);
257            for value in values {
258                self.print(separator);
259                self.print(value);
260            }
261        }
262        self.println("");
263    }
264
265    /// `values` の要素を区切りなしで連結し、末尾に改行を付けて出力します。
266    pub fn join_nospace<T: Display>(&mut self, values: impl IntoIterator<Item = T>) {
267        self.join(values, "");
268    }
269
270    /// `values` の要素を ASCII 空白1文字で区切り、末尾に改行を付けて出力します。
271    pub fn join_whitespace<T: Display>(&mut self, values: impl IntoIterator<Item = T>) {
272        self.join(values, " ");
273    }
274
275    /// `values` の要素を改行で区切って出力し、最後にも改行を付けます。
276    pub fn join_line<T: Display>(&mut self, values: impl IntoIterator<Item = T>) {
277        self.join(values, "\n");
278    }
279
280    /// バッファの内容を出力先へ書き出します。
281    ///
282    /// 書き込みに失敗した場合は panic します。
283    pub fn flush(&mut self) {
284        self.writer.flush().unwrap();
285    }
286}
287
288impl Writer<io::StdoutLock<'static>> {
289    /// 標準出力を出力先として作ります。
290    ///
291    /// 標準出力をロックしてラップします。出力を明示的に確定したい場合は [`flush`](Self::flush) を呼びます。
292    pub fn new() -> Self {
293        Self::with_writer(io::stdout().lock())
294    }
295}
296
297impl Default for Writer<io::StdoutLock<'static>> {
298    fn default() -> Self {
299        Self::new()
300    }
301}
302
303impl<W: Write> Drop for Writer<W> {
304    fn drop(&mut self) {
305        self.flush();
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::{Scanner, Writer};
312
313    #[test]
314    fn scans_supported_values() {
315        let mut scanner = Scanner::new("3 hello\n".as_bytes());
316        assert_eq!(scanner.token::<usize>(), 3);
317        assert_eq!(scanner.token::<String>(), "hello");
318    }
319
320    #[test]
321    fn writes_joined_values() {
322        let mut output = Vec::new();
323        {
324            let mut writer = Writer::with_writer(&mut output);
325            writer.join_whitespace([1, 2, 3]);
326        }
327        assert_eq!(output, b"1 2 3\n");
328    }
329}
atcoder/lib/io.rs

atcoder/lib/
io.rs
