AtCoder Beginner Contest 395 D

https://atcoder.jp/contests/abc395/tasks/abc395_d

鳩と巣を直接結び付けると2の処理に時間がかかるので、巣じゃなくてかごを用意して、1は鳩とかごの処理に、2はかごと巣の処理にすればよいです。

// Pigeon Swap
#![allow(non_snake_case)]

use std::collections::HashSet;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Query ////////////////////

enum Query {
    Move(usize, usize),
    Swap(usize, usize),
    Print(usize),
}

impl Query {
    fn read() -> Query {
        let v: Vec<usize> = read_vec();
        if v[0] == 1 {
            Query::Move(v[1]-1, v[2]-1)
        }
        else if v[0] == 2 {
            Query::Swap(v[1]-1, v[2]-1)
        }
        else {
            Query::Print(v[1]-1)
        }
    }
}


//////////////////// Nests ////////////////////

struct Nests {
    PtoB: Vec<usize>,           // 鳩 ⇒ かご
    BtoP: Vec<HashSet<usize>>,  // かご ⇒ 鳩
    BtoN: Vec<usize>,           // かご ⇒ 巣
    NtoB: Vec<usize>,           // 巣 ⇒ かご
}

impl Nests {
    fn relocate(&mut self, p: usize, n: usize) {
        let b_from = self.PtoB[p];  // かご
        let b_to = self.NtoB[n];    // かご
        if self.BtoP[b_from].remove(&p) {
            self.BtoP[b_to].insert(p);
            self.PtoB[p] = b_to
        }
    }
    
    fn swap(&mut self, n1: usize, n2: usize) {
        let b1 = self.NtoB[n1];
        let b2 = self.NtoB[n2];
        self.NtoB[n1] = b2;
        self.NtoB[n2] = b1;
        self.BtoN[b1] = n2;
        self.BtoN[b2] = n1
    }
    
    fn get_nest(&self, a: usize) -> usize {
        self.BtoN[self.PtoB[a]]
    }
    
    fn create(N: usize) -> Nests {
        let PtoB: Vec<usize> = (0..N).map(|n| n).collect();
        let BtoP: Vec<HashSet<usize>> = (0..N).map(|n|
                                                vec![n].into_iter().collect()).
                                                collect();
        let BtoN = PtoB.clone();
        let NtoB = PtoB.clone();
        Nests { PtoB, BtoP, BtoN, NtoB }
    }
}


//////////////////// process ////////////////////

fn read_input() -> (usize, usize) {
    let v: Vec<usize> = read_vec();
    let (N, Q) = (v[0], v[1]);
    (N, Q)
}

fn query(q: Query, nests: &mut Nests) {
    match q {
        Query::Move(a, b) => nests.relocate(a, b),
        Query::Swap(a, b) => nests.swap(a, b),
        Query::Print(a)   => println!("{}", nests.get_nest(a) + 1)
    }
}

fn F(N: usize, Q: usize) {
    let mut nests: Nests = Nests::create(N);
    for _ in 0..Q {
        let q = Query::read();
        query(q, &mut nests)
    }
}

fn main() {
    let (N, Q) = read_input();
    F(N, Q)
}