Back to Solutions

Collecting Numbers 2

Explanation

This problem is an extension of Collecting Numbers. Here, we are given m operations after the array, where elements at positions a and b are swapped and we need to find the number of trips after each operation.

This can be done by first calculating the count of the entire array before any operation is applied (check previous problem). Then for each operation, get the count of the trips added for pairs [(num1 - 1, num1), (num1, num1 + 1), (num2 - 1, num2), (num2, num2 + 1)].

Then swap the 2 numbers. And get the count of the trips for those pairs again. Subtrack the previous count from the new count and add it to the answer.

Unfortunately, this approach is slow and it results in a TLE verdict for Python. However, it passes for Rust and C++ (barely).
I'll try my best to come up with a better approach and update this post soon.

Code

use std::collections::{HashMap, HashSet};
use std::io;

fn main() {
    let inp = scan_vec::<usize>();
    let (_, m) = (inp[0], inp[1]);
    let mut arr = scan_vec::<usize>();

    let mut indices = HashMap::new();
    for (i, e) in arr.iter().enumerate() {
        indices.insert(*e, i);
    }

    let mut count = 1;
    for i in 1..arr.len() {
        if indices[&i] > indices[&(i + 1)] {
            count += 1;
        }
    }

    for _ in 0..m {
        let inp = scan_vec::<usize>();
        let (a, b) = (inp[0], inp[1]);
        let (x, y) = (arr[a - 1], arr[b - 1]);
        arr[a - 1] = y;
        arr[b - 1] = x;

        let prev = get_partial_count(&arr, &indices, x, y);

        let (a, b) = (indices[&x], indices[&y]);
        *indices.get_mut(&x).unwrap() = b;
        *indices.get_mut(&y).unwrap() = a;

        let curr = get_partial_count(&arr, &indices, x, y);

        count += curr - prev;
        println!("{}", count);
    }
}

fn get_partial_count(arr: &Vec<usize>, indices: &HashMap<usize, usize>, num1: usize, num2: usize) -> usize {
    let mut count = 0;
    let mut counted = HashSet::new();
    for i in num1 - 1..num1 + 1 {
        if i == 0 || i >= arr.len() {
            continue;
        }
        counted.insert((i, i + 1));
        if indices[&i] > indices[&(i + 1)] {
            count += 1;
        }
    }

    for i in num2 - 1..num2 + 1 {
        if i == 0 || i >= arr.len() || counted.contains(&(i, i + 1)) {
            continue;
        }
        if indices[&i] > indices[&(i + 1)] {
            count += 1;
        }
    }

    count
}

fn scan_vec<T: std::str::FromStr>() -> Vec<T> {
    let mut inp = String::new();
    io::stdin().read_line(&mut inp).unwrap();
    inp.trim()
        .split_whitespace()
        .map(|e| e.parse().ok().unwrap())
        .collect()
}