AtCoder Beginner Contest 432 E

https://atcoder.jp/contests/abc432/tasks/abc432_e

和の中の式は、なら常に、なら3つの範囲に分けて、なら、なら、ならとなります。
値とそれがいくつあるかを木にすればよいです。

// Clamp
#![allow(non_snake_case)]


//////////////////// 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()
}


//////////////////// Segment Tree ////////////////////

type Node = (usize, usize);     // (個数, 和)
type Range = (usize, usize);
type NodeRange = (usize, usize, usize);     // (first, last, index)

struct SegmentTree {
    n: usize,
    v: Vec<Node>,
}

impl SegmentTree {
    fn ceil_two_pow(n: usize) -> usize {
        if n == 1 { 1 } else { SegmentTree::ceil_two_pow((n+1)/2) * 2 }
    }
    
    fn new(B: Vec<usize>) -> SegmentTree {
        let m = B.len();
        let n = SegmentTree::ceil_two_pow(m);
        let mut v: Vec<Node> = vec![(0, 0); n*2-1];
        for i in 0..m {
            v[n-1+i] = (B[i], i * B[i])
        }
        for i in (0..n-1).rev() {
            let (n1, s1) = v[i*2+1];
            let (n2, s2) = v[i*2+2];
            v[i] = (n1+n2, s1+s2)
        }
        SegmentTree { n, v }
    }
    
    fn is_leaf(&self, i: usize) -> bool {
        i >= self.n - 1
    }
    
    fn update(&mut self, k: usize, node: Node, (first, last, n): NodeRange) {
        if self.is_leaf(n) {
            self.v[n] = node
        }
        else {
            let mid = (first + last) / 2;
            if k < mid {
                self.update(k, node, (first, mid, n*2+1))
            }
            else {
                self.update(k, node, (mid, last, n*2+2))
            }
            let (n1, s1) = self.v[n*2+1];
            let (n2, s2) = self.v[n*2+2];
            self.v[n] = (n1+n2, s1+s2)
        }
    }
    
    fn increase(&mut self, y: usize) {
        let (n, s) = self.v[self.n-1+y];
        self.update(y, (n+1, s+y), (0, self.n, 0));
    }
    
    fn decrease(&mut self, y: usize) {
        let (n, s) = self.v[self.n-1+y];
        self.update(y, (n-1, s-y), (0, self.n, 0));
    }
    
    fn count_core(&self, rng: Range,
                        (first, last, n): NodeRange) -> (usize, usize) {
        if rng.0 <= first && last <= rng.1 {
            self.v[n]
        }
        else {
            let mid = (first + last) / 2;
            let mut m: usize = 0;
            let mut s: usize = 0;
            if rng.0 < mid {
                let (n1, s1) = self.count_core(rng, (first, mid, n*2+1));
                m += n1;
                s += s1;
            }
            if mid < rng.1 {
                let (n2, s2) = self.count_core(rng, (mid, last, n*2+2));
                m += n2;
                s += s2;
            }
            (m, s)
        }
    }
    
    fn count(&self, first: usize, last: usize) -> (usize, usize) {
        self.count_core((first, last), (0, self.n, 0))
    }
}


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

enum Query {
    Change(usize, usize),
    Sum(usize, usize),
}

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

fn max_value(qs: &Vec<Query>) -> usize {
    let mut max_y: usize = 0;
    for q in qs {
        if let Query::Change(_, y) = *q {
            max_y = max(max_y, y)
        }
    }
    max_y
}


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

use std::cmp::{min, max};

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

fn frequency(A: &Vec<usize>, upper: usize) -> Vec<usize> {
    let mut B: Vec<usize> = vec![0; upper+1];
    for &a in A {
        B[a as usize] += 1
    }
    B
}

fn change(x: usize, y: usize, tree: &mut SegmentTree, A: &mut Vec<usize>) {
    tree.decrease(A[x]);
    A[x] = y;
    tree.increase(A[x])
}

fn sum(l_: usize, r_: usize, tree: &SegmentTree) -> usize {
    let l = min(l_, tree.n);
    let r = min(r_, tree.n);
    if l >= r {
        let (n, _) = tree.v[0];
        n * l
    }
    else {
        let (n1,  _) = tree.count(0, l);
        let ( _, s2) = tree.count(l, r);
        let (n3,  _) = tree.count(r, tree.n);
        n1 * l + s2 + n3 * r
    }
}

fn query(q: Query, tree: &mut SegmentTree, A: &mut Vec<usize>) {
    match q {
        Query::Change(x, y) => change(x, y, tree, A),
        Query::Sum(l, r)    => println!("{}", sum(l, r, tree)),
    }
}

fn F(Q: usize, mut A: Vec<usize>) {
    let queries: Vec<Query> = (0..Q).map(|_| Query::read()).collect();
    let upper = max(A.iter().cloned().max().unwrap(), max_value(&queries));
    let B = frequency(&A, upper);
    let mut tree = SegmentTree::new(B);
    for q in queries {
        query(q, &mut tree, &mut A)
    }
}

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