AtCoder Beginner Contest 452 F

https://atcoder.jp/contests/abc452/tasks/abc452_f

転倒数は1個右に追加したり左から除いたりしたときの増減を求めるのは簡単です。右に追加するときはあらかじめ木を作っておいて、自分より大きいものの数を数えればいいです。そうやって尺取り法のように木を更新しながらカウントしていきます。ただし、その周辺でずっとKになるところではややこしいです。

// Interval Inversion Count
#![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()
}


//////////////////// 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;     // ノードがいくつあるか
    fn op(n1: &Self::S, n2: &Self::S) -> Self::S {
        n1 + n2
    }
    fn e() -> Self::S { 0 }
}


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

struct Tree {
    seg: SegTree::<Sum>
}

impl Tree {
    fn add(&mut self, i: usize) {
        self.seg.set(i, 1)
    }
    
    fn remove(&mut self, i: usize) {
        self.seg.set(i, 0)
    }
    
    fn count(&self, first: usize, last: usize) -> usize {
        self.seg.prod(first, last)
    }
    
    fn new(N: usize) -> Tree {
        let seg = SegTree::new(N);
        Tree { seg }
    }
}


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

fn read_input() -> (usize, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let K = v[1];
    let P: Vec<usize> = read_vec::<usize>().into_iter().
                                            map(|i| i-1).collect();
    (K, P)
}

// どれだけlastを大きくしても転倒数は変わらないか
fn num_inc_last(last: usize, P: &Vec<usize>, tree: &Tree) -> usize {
    let N = P.len();
    if last == N || tree.count(P[last], N) != 0 {
        return 1
    }
    
    let mut p = last + 1;
    while p < N && P[p-1] < P[p] {
        p += 1
    }
    p + 1 - last
}

fn num_inc_first(first: usize, P: &Vec<usize>, tree: &Tree) -> usize {
    let last = first + tree.seg.data[0];
    for p in first..last {
        if tree.count(0, P[p]) != p - first {
            return p + 1 - first
        }
    }
    last - first
}

fn F(K: usize, P: Vec<usize>) -> usize {
    let N = P.len();
    let mut tree = Tree::new(N);
    let mut first: usize = 0;
    let mut last: usize = 1;
    tree.add(P[first]);
    let mut counter: usize = 0;
    let mut n: usize = 0;
    while last <= N {
        if n == K {
            let n1 = num_inc_last(last, &P, &tree);
            for p in last..(last+n1-1) {
                tree.add(P[p])
            }
            let mut n2 = num_inc_first(first, &P, &tree);
            counter += n1 * n2;
            if first + n2 > last {
                // first >= lastのときもカウントしているのでその分引く
                let d = first + n1 - last;
                counter -= d * (d + 1) / 2
            }
            for p in last+n1-1..(last+n1).min(N) {
                tree.add(P[p]);
                n += tree.count(P[p]+1, N)
            }
            for p in first..first+n2 {
                tree.remove(P[p]);
                n -= tree.count(0, P[p])
            }
            last += n1;
            first += n2
        }
        else if n < K {
            if last < N {
                n += tree.count(P[last]+1, N);
                tree.add(P[last]);
            }
            last += 1
        }
        else {
            n -= tree.count(0, P[first]);
            tree.remove(P[first]);
            first += 1
        }
    }
    counter
}

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