AtCoder Beginner Contest 441 E

https://atcoder.jp/contests/abc441/tasks/abc441_e

場合の数を数えるから分割統治法かなと思って考えても、各領域でAとBの数の差ごとに場合の数を考えないといけなくて、領域の結合でさらにAとBの差ごとの場合の数を数えるのにかかってしまって時間が間に合いません。
Aを1、Bを-1、Cを0として累積和を考えると、入力例1は[0, 1, 1, 0, -1, -1, 0, -1, -1, 0, -1]となります。そうすると、最初の要素は0でそれより大きい個数を数えます。いつものように木を作って、ある要素より大きい要素の個数をカウントして、その要素をデクリメントします。

// E - A > B substring
#![allow(non_snake_case)]

use std::collections::BTreeMap;


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


//////////////////// SegTree ////////////////////

trait Monoid {
    type S: Clone;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn e() -> Self::S;
}

struct SegTree<M: Monoid> {
    n: usize,
    size: usize,
    data: Vec<M::S>,
}

impl<M: Monoid> SegTree<M> {
    fn is_leaf(&self, i: usize) -> bool {
        i >= self.size - 1
    }
    
    fn new(n: usize) -> Self {
        let size = n.next_power_of_two();
        SegTree {
            n,
            size,
            data: vec![M::e(); size*2-1],
        }
    }
    
    fn from(v: Vec<M::S>) -> Self {
        let n = v.len();
        let size = n.next_power_of_two();
        let mut data = vec![M::e(); size*2-1];
        for i in 0..n {
            data[i+size-1] = v[i].clone();
        }
        let mut seg = SegTree { n, size, data };
        for i in (0..size-1).rev() {
            seg.data[i] = M::op(&seg.data[2*i+1], &seg.data[2*i+2]);
        }
        seg
    }
    
    fn set(&mut self, mut i: usize, x: M::S) {
        i = i + self.size - 1;
        self.data[i] = x;
        while i != 0 {
            i = (i - 1) / 2;
            self.data[i] = M::op(&self.data[2*i+1], &self.data[2*i+2]);
        }
    }
    
    fn prod(&self, mut l: usize, mut r: usize) -> M::S {
        let mut sml = M::e();
        let mut smr = M::e();
        l = l + self.size - 1;
        r = r + self.size - 1;
        while l < r {
            if l % 2 == 0 {
                sml = M::op(&sml, &self.data[l]);
                l += 1;
            }
            if r % 2 == 0 {
                r -= 1;
                smr = M::op(&self.data[r], &smr);
            }
            l = (l - 1) / 2;
            r = (r - 1) / 2
        }
        M::op(&sml, &smr)
    }
    
    fn get(&self, i: usize) -> M::S {
        self.data[i+self.size-1].clone()
    }
}


//////////////////// Monoid ////////////////////

struct Sum;
impl Monoid for Sum {
    type S = (usize, i32);  // (個数, 値)
    fn op((n1, e1): &Self::S, (n2, e2): &Self::S) -> Self::S {
        (n1+n2, *e2)
    }
    fn e() -> Self::S { (0, i32::MAX) }
}


//////////////////// Tree ////////////////////

struct Tree {
    seg: SegTree::<Sum>,
    B: Vec<i32>,    // 値を保持
}

impl Tree {
    fn size(&self) -> usize {
        self.B.len()
    }
    
    // 与えられた値未満のBの最大のindex
    fn find_B_index(&self, b: i32) -> usize {
        let mut first: usize = 0;
        let mut last: usize = self.B.len();
        while last - first > 1 {
            let mid = (first + last) / 2;
            if self.B[mid] >= b {
                last = mid
            }
            else {
                first = mid
            }
        }
        first
    }
    
    fn decrement(&mut self, b: i32) {
        let i = if b == self.B[0] { 0 } else { self.find_B_index(b) + 1 };
        let (n, e) = self.seg.get(i);
        self.seg.set(i, (n-1, e))
    }
    
    // 与えられた値より大きいBの個数の和
    fn sum(&self, b: i32) -> usize {
        let i = if b == self.B[0] { 0 } else { self.find_B_index(b) + 1 };
        self.seg.prod(i+1, self.seg.n).0
    }
    
    fn from(A: Vec<(usize, i32)>) -> Tree {
        let B: Vec<i32> = A.iter().map(|&(n, e)| e).collect();
        let seg = SegTree::from(A);
        Tree { seg, B }
    }
}


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

fn read_input() -> String {
    let _N: usize = read();
    let S: String = read();
    S
}

fn accumulate(S: String) -> Vec<i32> {
    let N = S.len();
    let mut acc: Vec<i32> = vec![0; N+1];
    for (i, c) in S.chars().enumerate() {
        let e: i32 = match c {
                        'A' =>  1,
                        'B' => -1,
                        _   =>  0
        };
        acc[i+1] = acc[i] + e
    }
    acc
}

fn frequency(acc: &Vec<i32>) -> Vec<(usize, i32)> {
    let mut m: BTreeMap<i32, usize> = BTreeMap::new();
    for &n in acc.iter() {
        let e = m.entry(n).or_insert(0);
        *e += 1
    }
    m.into_iter().map(|(e, n)| (n, e)).collect::<Vec<_>>()
}

fn F(S: String) -> usize {
    let N = S.len();
    let acc = accumulate(S);
    let freq = frequency(&acc);
    let mut tree = Tree::from(freq);
    let mut counter: usize = 0;
    for e in acc {
        counter += tree.sum(e);
        tree.decrement(e)
    }
    counter
}

fn main() {
    let S = read_input();
    println!("{}", F(S))
}