1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
//!
//! Binary Indexed Tree (BIT)
//!
//! 参考:<https://algo-logic.info/binary-indexed-tree/>
//!

///
/// Binary Indexed Tree (BIT)
///
pub struct BinaryIndexedTree {
    n: usize,
    bit: Vec<isize>,
}

impl BinaryIndexedTree {
    pub fn new(n: usize) -> BinaryIndexedTree {
        let bit = vec![0; n + 1];

        BinaryIndexedTree { n: n + 1, bit }
    }

    pub fn add(&mut self, i: usize, x: isize) {
        assert_ne!(i, 0, "i is 1-index");

        let mut index = i as isize;

        while (index as usize) < self.n {
            self.bit[index as usize] += x;

            // i & -i => i の最後の1のビット
            index += index & -index;
        }
    }

    pub fn sum(&self, i: usize) -> isize {
        assert_ne!(i, 0, "i is 1-index");

        let mut s = 0;

        let mut index = i as isize;

        while index > 0 {
            s += self.bit[index as usize];

            // i & -i => i の最後の1のビット
            index -= index & -index;
        }

        s
    }
}

#[cfg(test)]
mod tests {
    use super::BinaryIndexedTree;

    #[test]
    fn test_bit() {
        let mut tree = BinaryIndexedTree::new(5);

        tree.add(1, 10);
        tree.add(2, 20);
        tree.add(3, 30);
        tree.add(4, 40);
        tree.add(5, 50);

        assert_eq!(tree.sum(1), 10);
        assert_eq!(tree.sum(2), 30);
        assert_eq!(tree.sum(3), 60);
        assert_eq!(tree.sum(4), 100);
        assert_eq!(tree.sum(5), 150);
    }

    #[test]
    fn test_inversion() {
        assert_eq!(inversion(5, vec![3, 1, 5, 4, 2]), 5);
        assert_eq!(inversion(6, vec![1, 2, 3, 4, 5, 6]), 0);
        assert_eq!(inversion(7, vec![7, 6, 5, 4, 3, 2, 1]), 21);
        assert_eq!(
            inversion(
                20,
                vec![19, 11, 10, 7, 8, 9, 17, 18, 20, 4, 3, 15, 16, 1, 5, 14, 6, 2, 13, 12]
            ),
            114
        );
    }

    fn inversion(n: usize, a: Vec<usize>) -> isize {
        // 転倒数を求める

        let mut bit = BinaryIndexedTree::new(n);

        let mut inversion = 0;

        for x in a {
            inversion += bit.sum(n) - bit.sum(x);

            bit.add(x, 1);
        }

        inversion
    }
}