< MizarさんのAI生成コード置き場

input_lite

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

// -*- coding:utf-8-unix -*-
pub use std::io::prelude::*;

/// 1 トークンを任意の型として読むためのマクロ群。
///
/// 使い方の例:
/// - `read_value!(b, usize)`       : 1 個の整数を読む
/// - `read_value!(b, (i64, i64))` : 2 個読んでタプルにする
/// - `read_value!(b, [u32; n])`   : n 個読んで `Vec<u32>` にする
#[macro_export]
macro_rules! read_value {
    // タプルを読む。
    // 例えば `(usize, i64)` なら 2 個のトークンを順に読んでタプル化する。
    ($b:expr, ( $($t:tt),* )) => {
        ( $(read_value!($b, $t)),* )
    };

    // 長さ付き配列風記法。
    // 実際には固定長配列ではなく Vec として返す。
    ($b:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($b, $t)).collect::<Vec<_>>()
    };

    // 1 トークンの先頭 1 byte を読む。
    // ASCII 文字を 1 文字読む用途を想定している。
    ($b:expr, byte) => {
        token($b, |s| s[0])
    };

    // 1 トークンをそのまま `Vec<u8>` として読む。
    ($b:expr, Bytes) => {
        token($b, |s| s.to_vec())
    };

    // 1 トークンを UTF-8 文字列とみなし、`Vec<char>` に変換する。
    ($b:expr, Chars) => {
        token($b, |s| std::str::from_utf8(s).unwrap().chars().collect::<Vec<_>>())
    };

    // 1-indexed な整数入力を 0-indexed に直すための定番記法。
    // 0 が来たら仕様違反として panic する。
    ($b:expr, usize1) => {
        read_value!($b, usize).checked_sub(1).unwrap()
    };
    ($b:expr, Usize1) => {
        read_value!($b, usize).checked_sub(1).unwrap()
    };
    ($b:expr, Isize1) => {
        read_value!($b, isize).checked_sub(1).unwrap()
    };

    // 一般の型。
    // 1 トークンを UTF-8 として解釈し、`parse()` で目的の型へ変換する。
    ($b:expr, $t:ty) => {
        token($b, |s| {
            std::str::from_utf8(s).unwrap().parse::<$t>().unwrap()
        })
    };
}

/// `input!` の内部で、`n: usize, mut a: [i64; n], ...` のような宣言列を順に処理する。
#[macro_export]
macro_rules! input_inner {
    // 引数が尽きたら終了。
    ($b:expr) => {};
    ($b:expr,) => {};

    // `mut x: T` の形。
    ($b:expr, mut $var:ident : $t:tt $($r:tt)*) => {
        let mut $var = read_value!($b, $t);
        input_inner!{$b $($r)*}
    };

    // `x: T` の形。
    ($b:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($b, $t);
        input_inner!{$b $($r)*}
    };
}

/// 競技プログラミング風の入力マクロ。
///
/// 例: 標準入力から読む場合
/// ```ignore
/// input! {
///     n: usize,
///     mut a: [i64; n],
///     s: Chars,
/// }
/// a.sort_unstable();
/// ```
///
/// `source = ...` を指定すると、その `BufRead` から読む。
///
/// 例: `&[u8]` (BufRead を実装) から読む場合
/// ```ignore
/// let mut src = &b"5 -2 -1 0 1 2 abc"[..];
/// input! {
///     source = &mut src,
///     n: usize,
///     mut a: [i64; n],
///     s: Chars,
/// }
/// a.sort_unstable();
/// ```
///
/// source を指定しない場合は標準入力を lock して読む。
#[macro_export]
macro_rules! input {
    // 明示的に入力元を与える場合。
    (source = $s:expr, $($r:tt)*) => {
        input_inner!{$s, $($r)*}
    };

    // 標準入力から読む場合。
    ($($r:tt)*) => {
        let mut i = std::io::stdin().lock();
        input_inner!{&mut i, $($r)*}
        std::mem::drop(i);
    };
}

/// `BufRead` から空白区切りで 1 トークン読み、
/// その byte 列に関数 `f` を適用して値を返す。
///
/// 実装方針:
/// - `fill_buf()` / `consume()` を直接使う
/// - 短いトークンはコピーせずに処理する
/// - バッファ境界をまたぐ長いトークンだけ `Vec<u8>` に退避する
///
/// 注意:
/// - 入力不足時は `assert!(!b.is_empty())` で panic する
/// - 区切り判定は `is_ascii_whitespace()` を使う
pub fn token<R: BufRead, T>(r: &mut R, f: impl Fn(&[u8]) -> T) -> T {
    // まず先頭の空白を飛ばし、最初の非空白文字から始まるバッファ断片を得る。
    //
    // 戻り値:
    // - `b`: 最初の非空白文字以降のスライス
    // - `l`: 元のバッファ先頭から最初の非空白までの長さ
    let (b, l) = loop {
        let b = r.fill_buf().unwrap();
        // トークンが必要なのに EOF なら、これは入力不足エラー。
        assert!(!b.is_empty(), "unexpected EOF while looking for token");

        // 非空白文字が見つかったら、そこから先をトークン候補とする。
        if let Some(p) = b.iter().position(|&c| !c.is_ascii_whitespace()) {
            break (&b[p..], p);
        }

        // 今回のバッファ断片がすべて空白なら、そのぶん全部捨てる。
        let l = b.len();
        r.consume(l);
    };

    // 同じバッファ断片の中でトークンが完結する場合。
    if let Some(p) = b.iter().position(|&c| c.is_ascii_whitespace()) {
        // `b[..p]` が 1 個のトークン。
        let t = &b[..p];
        let x = f(t);

        // 消費量は
        // - 先頭の空白 `l`
        // - トークン本体 `p`
        // - 直後の空白 1 文字
        //
        // 残りの空白が複数ある場合は、次回呼び出し時に先頭空白として処理される。
        r.consume(l + p + 1);
        return x;
    }

    // ここに来るのは、`b` の終端まで空白がなく、
    // トークンがバッファ境界をまたいで続いている場合。
    //
    // いったん現在の断片を Vec に退避し、後続の断片を継ぎ足していく。
    let mut t = b.to_vec();
    r.consume(l + t.len());

    while let Ok(b) = r.fill_buf() {
        // EOF ならトークン完了。
        if b.is_empty() {
            break; // EOF
        }

        // この断片の中で初めて空白が出たなら、そこまででトークン完了。
        if let Some(p) = b.iter().position(|&c| c.is_ascii_whitespace()) {
            t.extend_from_slice(&b[..p]);
            r.consume(p + 1); // トークン本体 + 直後の空白 1 文字
            break;
        }

        // まだ空白がないなら、断片全体がトークンの続き。
        t.extend_from_slice(b);
        let l = b.len();
        r.consume(l);
    }

    f(&t)
}