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 self::bitset::BitSet;
use std::cmp;
use std::collections::BTreeMap;

fn main() {
    let s = std::io::stdin();
    let mut sc = Scanner { stdin: s.lock() };

    let n = sc.read();
    let r: usize = sc.read();
    let x: Vec<usize> = sc.vec::<usize>(n).into_iter().map(|x| x - 1).collect();
    let mut used = vec![0; n];
    let mut phase = 1;
    for i in 0..n {
        if used[i] != 0 {
            continue;
        }
        let mut cur = i;
        while used[cur] == 0 {
            used[cur] = phase;
            cur = x[cur];
        }
        phase += 1;
    }

    let mut map = BTreeMap::new();
    for p in used.into_iter() {
        *map.entry(p).or_insert(0) += 1;
    }

    let mut count = vec![0; n + 1];
    for (_, c) in map.into_iter() {
        count[c] += 1;
    }

    let one = count[1];
    let count = count
        .into_iter()
        .enumerate()
        .filter(|&(v, c)| v > 1 && c > 0)
        .collect::<Vec<(usize, usize)>>();

    let r = cmp::min(r, n - r);
    let mut b = BitSet::new(1);
    b.set(0, true);
    for (v, count) in count.into_iter() {
        for _ in 0..count {
            let c = b.shift_left(v);
            b |= c;
        }
    }

    for i in 0..(one + 1) {
        if r >= i && b.get(r - i) {
            println!("Yes");
            return;
        }
    }

    println!("No");
}

pub mod bitset {
    use std::ops::{BitOrAssign, Shl};
    const ONE_VALUE_LENGTH: usize = 63;
    const MAXIMUM: u64 = (1u64 << ONE_VALUE_LENGTH as u64) - 1;

    pub fn get_bit_position(index: usize) -> (usize, usize) {
        let data_index = index / ONE_VALUE_LENGTH;
        let bit_index = index % ONE_VALUE_LENGTH;
        (data_index, bit_index)
    }

    #[derive(PartialEq, Clone, Debug)]
    pub struct BitSet {
        data: Vec<u64>,
    }

    impl BitOrAssign for BitSet {
        fn bitor_assign(&mut self, rhs: Self) {
            if self.data.len() < rhs.data.len() {
                self.data.resize(rhs.data.len(), 0);
            }
            let n = if self.data.len() > rhs.data.len() {
                rhs.data.len()
            } else {
                self.data.len()
            };
            for i in 0..n {
                assert!(self.data[i] <= MAXIMUM);
                assert!(rhs.data[i] <= MAXIMUM);
                self.data[i] |= rhs.data[i];
            }
        }
    }

    impl Shl<usize> for BitSet {
        type Output = Self;
        fn shl(self, rhs: usize) -> Self {
            self.shift_left(rhs)
        }
    }

    impl BitSet {
        pub fn new(n: usize) -> Self {
            let size = (n + ONE_VALUE_LENGTH - 1) / ONE_VALUE_LENGTH;
            BitSet {
                data: vec![0; size],
            }
        }

        pub fn new_from(value: u64) -> Self {
            BitSet { data: vec![value] }
        }

        pub fn set(&mut self, index: usize, value: bool) {
            let (data_index, bit_index) = get_bit_position(index);
            assert!(self.data.len() > data_index);
            if value {
                self.data[data_index] |= (1 << bit_index) as u64;
            } else {
                let tmp = MAXIMUM ^ (1 << bit_index) as u64;
                self.data[data_index] &= tmp;
            }
        }

        pub fn get(&mut self, index: usize) -> bool {
            let (data_index, bit_index) = get_bit_position(index);
            self.data.len() > data_index && self.data[data_index] & (1u64 << bit_index as u64) != 0
        }

        pub fn shift_left(&self, shift: usize) -> Self {
            let mut next_data = Vec::new();
            let prefix_empty_count = shift / ONE_VALUE_LENGTH;
            let shift_count = (shift % ONE_VALUE_LENGTH) as u64;
            for _ in 0..prefix_empty_count {
                next_data.push(0);
            }

            let mut from_previous = 0;
            let room = ONE_VALUE_LENGTH as u64 - shift_count;
            for &data in self.data.iter() {
                let overflow = (data >> room) << room;
                let rest = data - overflow;
                let value = (rest << shift_count) + from_previous;
                assert!(value <= MAXIMUM);
                next_data.push(value);
                from_previous = overflow >> room;
            }
            if from_previous > 0 {
                next_data.push(from_previous);
            }
            BitSet { data: next_data }
        }
    }
}

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()
    }
}