Cod sursa(job #3311968)

Utilizator sad_carrotVisan Sebastian sad_carrot Data 25 septembrie 2025 09:52:19
Problema Datorii Scor 100
Compilator rs Status done
Runda Arhiva de probleme Marime 3.87 kb
use std::{
    cmp::Ordering::{Equal, Greater, Less},
    fs,
};

struct Node {
    value: i64,
    sum: i64, // of the subtree
    pos: usize,
    left: Subtree,
    right: Subtree,
}

struct Subtree(Option<Box<Node>>);

impl Subtree {
    pub fn new(nums: &[i64], before: usize) -> Self {
        if nums.is_empty() {
            Subtree(None)
        } else {
            let middle = nums.len() / 2;
            Subtree(Some(Box::new(Node {
                value: nums[middle],
                sum: nums[middle],
                pos: before + middle,
                left: Subtree::new(&nums[..middle], before),
                right: Subtree::new(&nums[middle + 1..], before + middle + 1),
            })))
        }
    }
    pub fn calculate_sum(&mut self) -> i64 {
        match &mut self.0 {
            None => 0,
            Some(node) => {
                let leftsum = node.left.calculate_sum();
                let rightsum = node.right.calculate_sum();
                node.sum = node.sum + leftsum + rightsum;
                node.sum
            }
        }
    }
    pub fn sum(&self) -> i64 {
        match &self.0 {
            None => 0,
            Some(node) => node.sum,
        }
    }
    pub fn left_query(&self, a: usize) -> i64 {
        match &self.0 {
            None => 0,
            Some(node) => match node.pos.cmp(&a) {
                Less => node.right.left_query(a),
                Greater => node.value + node.right.sum() + node.left.left_query(a),
                Equal => node.value + node.right.sum(),
            },
        }
    }
    pub fn right_query(&self, b: usize) -> i64 {
        match &self.0 {
            None => 0,
            Some(node) => match node.pos.cmp(&b) {
                Less => node.value + node.left.sum() + node.right.right_query(b),
                Greater => node.left.right_query(b),
                Equal => node.value + node.left.sum(),
            },
        }
    }
    pub fn query(&self, a: usize, b: usize) -> i64 {
        match &self.0 {
            None => 0,
            Some(node) => match node.pos.cmp(&a) {
                Less => node.right.query(a, b),
                Equal => match node.pos.cmp(&b) {
                    Equal => node.value,
                    _ => node.value + node.right.right_query(b),
                },
                Greater => match node.pos.cmp(&b) {
                    Greater => node.left.query(a, b),
                    Equal => node.value + node.left.left_query(a),
                    Less => node.value + node.left.left_query(a) + node.right.right_query(b),
                },
            },
        }
    }
    pub fn update(&mut self, pos: usize, val: i64) {
        match &mut self.0 {
            None => {}
            Some(node) => {
                match node.pos.cmp(&pos) {
                    Equal => {
                        node.value -= val;
                        node.sum = val
                    }
                    Less => node.right.update(pos, val),
                    Greater => node.left.update(pos, val),
                }
                node.sum = node.value + node.left.sum() + node.right.sum();
            }
        }
    }
}

fn main() {
    let input = fs::read_to_string("datorii.in").unwrap();
    let content = input
        .split_ascii_whitespace()
        .map(|s| s.parse::<i64>().unwrap())
        .collect::<Vec<i64>>();
    let n = content[0] as usize;
    let nums = &content[2..n + 2];
    let queries = &content[n + 2..];

    let mut arbint = Subtree::new(nums, 0);
    arbint.calculate_sum();

    let output: String = queries
        .chunks(3)
        .map(|q| {
            if q[0] == 0 {
                arbint.update(q[1] as usize - 1, q[2]);
                "".to_string()
            } else {
                format!("{}\n", arbint.query(q[1] as usize - 1, q[2] as usize - 1))
            }
        })
        .collect();
    fs::write("datorii.out", output).unwrap();
}