Problem:
D: Two Colors Sort

問題

umg 君は散歩中に 1,2, ... ,N を並び変えてできる長さ N の数列 P_1, P_2, ..., P_N を見つけました。

umg 君は不思議な力を使うことで、同じ色に塗られた異なる数を二つ選んで場所を交換することができます。

umg 君は、数列に含まれる数のうち R 個を赤に、残りの N-R 個を青に塗ることで数列を昇順に並び替えられるようにしたいと考えました。

umg 君が目標を達成できるかどうか判定してください。

ただし、数字はとても重いので不思議な力を使わずに動かすことはできません。また、umg 君は天才なので不思議な力を任意の回数使うことができます。

入力形式

N R
P_1 P_2 ... P_N

制約

1 \leq N \leq 3 \times 10^5

1 \leq R \leq N

1\leq P_i \leq N

P_i \neq P_j (1 \leq i < j \leq N)

入力は全て整数である。

出力形式

umg 君が目的を達成できるなら Yes、そうでなければ No を一行に出力せよ。

入力例 1

3 2
1 3 2

出力例 1

Yes

1 を青に、2,3 を赤に塗ることで目標が達成できます。

入力例 2

5 2
1 2 3 4 5

出力例 2

Yes

初めから昇順に並んでいます。

入力例 3

10 7
3 4 8 5 7 6 2 10 1 9

出力例 3

No

Solution in Rust (judge verdict: Accepted)

use std::cmp;
use std::collections::BTreeMap;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };
    let n: usize = sc.read();
    let r: usize = sc.read();
    let p: Vec<usize> = sc.vec::<usize>(n).into_iter().map(|x| x - 1).collect();

    let mut used = vec![0; n];
    let mut cur = 1;
    for i in 0..n {
        if used[i] != 0 {
            continue;
        }
        let mut pos = i;
        while used[pos] == 0 {
            used[pos] = cur;
            pos = p[pos];
        }
        cur += 1;
    }

    let count = used
        .into_iter()
        .fold(BTreeMap::new(), |mut map, x| {
            *map.entry(x).or_insert(0) += 1;
            map
        })
        .into_iter()
        .fold(BTreeMap::new(), |mut map, (_, count)| {
            *map.entry(count).or_insert(0) += 1;
            map
        });

    let r = cmp::min(r, n - r);
    let mut dp = vec![-1; r + 1];
    dp[0] = 0;
    for (a, m) in count.into_iter() {
        for j in 0..(r + 1) {
            if dp[j] >= 0 {
                dp[j] = m;
            } else if j < a || dp[j - a] <= 0 {
                dp[j] = -1;
            } else {
                dp[j] = dp[j - a] - 1;
            }
        }
    }

    if dp[r] >= 0 {
        println!("Yes");
    } else {
        println!("No");
    }
}

pub struct Scanner<R> {
    stdin: R,
}

impl<R: std::io::Read> Scanner<R> {
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .stdin
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n')
            .take_while(|&b| b != b' ' && b != b'\n')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }
}