Cod sursa(job #2830133)

Utilizator bcebereBogdan Cebere bcebere Data 9 ianuarie 2022 15:58:28
Problema Cautare binara Scor 40
Compilator rs Status done
Runda Arhiva educationala Marime 4.27 kb
use std::fs::File;
use std::io::prelude::*;
use std::io::LineWriter;
use std::io::{self, BufRead};

struct Query {
    qtype: u32,
    qval: i32,
}

fn read_data(path: &str) -> Result<(usize, Vec<i32>, Vec<Query>), 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 testcases = lines.next().unwrap()?.parse::<usize>().unwrap();

    let mut tests: Vec<Query> = Vec::with_capacity(testcases);
    for line in lines {
        let vals = line?;
        let mut tokens = vals.split(" ");

        tests.push(Query {
            qtype: tokens.next().unwrap().parse::<u32>().unwrap(),
            qval: tokens.next().unwrap().parse::<i32>().unwrap(),
        });
    }
    Ok((n, data, tests))
}

fn write_data(path: &str, data: Vec<String>) -> Result<(), std::io::Error> {
    let file = File::create(path)?;
    let mut file = LineWriter::new(file);

    for item in data {
        file.write_all((item + "\n").as_bytes())?;
    }
    Ok(())
}

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, tests) = read_data("cautbin.in").unwrap();
    assert!(n == data.len());

    let mut results: Vec<String> = Vec::with_capacity(tests.len());

    for query in tests {
        let mut res = match query.qtype {
            0 => search_exact(&data, query.qval),
            1 => search_exact_or_smaller(&data, query.qval) as isize,
            2 => search_exact_or_greater(&data, query.qval) as isize,
            _ => panic!("invalid query"),
        };
        if res >= 0 {
            res += 1;
        }
        results.push(res.to_string());
    }
    write_data("cautbin.out", results).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);
    }
}