Cod sursa(job #2830175)

Utilizator bcebereBogdan Cebere bcebere Data 9 ianuarie 2022 16:43:52
Problema Cautare binara Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 4.03 kb
use std::fs::File;
use std::io::prelude::*;
use std::io::{self, BufRead};
use std::io::BufReader;
use std::io::BufWriter;

fn read_data(path: &str) -> Result<(usize, Vec<i32>,  std::io::Lines<BufReader<File>>), std::io::Error> {
    let file = File::open(path)?;
    let mut lines = io::BufReader::new(file).lines();

    let n = lines.next().unwrap()?.parse::<usize>().unwrap();

    let mut data: Vec<i32> = Vec::with_capacity(n);
    for raw_item in lines.next().unwrap()?.split(" ") {
        if let Ok(val) = raw_item.parse::<i32>() {
            data.push(val);
        }
    }

    let _ = lines.next().unwrap()?.parse::<usize>().unwrap();
    Ok((n, data, lines))
}

fn data_writer(path: &str) -> Result<BufWriter<File>, std::io::Error> {
    let file = File::create(path)?;
    let stream = BufWriter::new(file);
    Ok(stream)
}

fn get_bits(data: &Vec<i32>) -> usize {
    let mut bits: usize = 0;
    while (1 << bits) <= data.len() {
        bits += 1;
    }

    bits
}

fn search_exact_or_smaller(data: &Vec<i32>, item: i32) -> usize {
    let mut final_pos: usize = 0;

    let bits = get_bits(data);

    for bit in (0..bits + 1).rev() {
        let tmp_final_pos = final_pos + (1 << bit);

        if tmp_final_pos >= data.len() {
            continue;
        }
        if data[tmp_final_pos] <= item {
            final_pos = tmp_final_pos;
        }
    }

    final_pos
}

fn search_exact_or_greater(data: &Vec<i32>, item: i32) -> usize {
    let mut final_pos: usize = data.len();

    let bits = get_bits(data);

    for bit in (0..bits).rev() {
        let tmp_final_pos = final_pos as isize - (1 << bit);

        if tmp_final_pos < 0 {
            continue;
        }
        if data[tmp_final_pos as usize] >= item {
            final_pos = tmp_final_pos as usize;
        }
    }

    final_pos
}

fn search_exact(data: &Vec<i32>, item: i32) -> isize {
    let pos = search_exact_or_smaller(data, item);

    match data[pos] == item {
        false => -1,
        true => pos as isize,
    }
}

fn main() {
    let (n, data, lines) = read_data("cautbin.in").unwrap();
    let mut writer = data_writer("cautbin.out").unwrap();

    assert!(n == data.len());

    for line in lines {
        let vals = line.unwrap();
        let mut tokens = vals.split(" ");

        let qtype = tokens.next().unwrap().parse::<u32>().unwrap();
        let qval= tokens.next().unwrap().parse::<i32>().unwrap();

        let mut res = match qtype {
            0 => search_exact(&data, qval),
            1 => search_exact_or_smaller(&data, qval) as isize,
            2 => search_exact_or_greater(&data, qval) as isize,
            _ => panic!("invalid query"),
        };
        if res >= 0 {
            res += 1;
        }
        writer.write((res.to_string() + "\n").as_bytes()).unwrap();
    }
    writer.flush().unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn search_exact_test() {
        assert_eq!(search_exact(&vec![1, 3, 3, 3, 3, 5], 1), 0);
        assert_eq!(search_exact(&vec![1, 3, 3, 3, 3, 5], 5), 5);
        assert_eq!(search_exact(&vec![1, 3, 3, 3, 3, 5], 3), 4);
        assert_eq!(search_exact(&vec![1, 3, 3, 3, 3, 5], 8), -1);
        assert_eq!(search_exact(&vec![1, 3, 3, 3, 3, 5], 4), -1);
        assert_eq!(search_exact(&vec![1, 3, 5, 6], 6), 3);
    }

    #[test]
    fn search_exact_or_smaller_test() {
        assert_eq!(search_exact_or_smaller(&vec![1, 3, 3, 3, 3, 5], 1), 0);
        assert_eq!(search_exact_or_smaller(&vec![1, 3, 3, 3, 3, 5], 5), 5);
        assert_eq!(search_exact_or_smaller(&vec![1, 3, 3, 3, 3, 5], 3), 4);
        assert_eq!(search_exact_or_smaller(&vec![1, 3, 3, 3, 3, 5], 8), 5);
        assert_eq!(search_exact_or_smaller(&vec![1, 3, 3, 3, 3, 5], 4), 4);
    }
    #[test]
    fn search_exact_or_greater_test() {
        assert_eq!(search_exact_or_greater(&vec![1, 3, 3, 3, 3, 5], 1), 0);
        assert_eq!(search_exact_or_greater(&vec![1, 3, 3, 3, 3, 5], 5), 5);
        assert_eq!(search_exact_or_greater(&vec![1, 3, 3, 3, 3, 5], 3), 1);
        assert_eq!(search_exact_or_greater(&vec![1, 3, 3, 3, 3, 5], 4), 5);
    }
}