树状数组板子

/*
 * @Author: CY__HHH
 * @Date: 2020-10-28 11:29:58
 * @LastEditTime: 2020-10-28 11:58:49
 */
#include <vector>
#include <iostream>
#include <iomanip>

class BIT{
    private:
        std::vector<int> my_arr;
    public:
        BIT(std::vector<int> arr){
            int len = (int)arr.size();
            my_arr = std::vector<int> (len+1,0);
            for(int idx=1; idx!=len + 1; ++idx)
                my_arr[idx] = arr[idx-1];
            for(int idx=1; idx!=len + 1; ++idx){
                int idx2 = idx + (idx & -idx);
                my_arr[idx2] += my_arr[idx];
            }
        }

        int prefix_query(int idx) const{
            int sum = 0;
            for(++idx; idx > 0; idx-=(idx & -idx))
                sum += my_arr[idx];
            return sum;
        }

        int range_query(int from_idx, int to_idx) const{
            if(from_idx == 0)
                return prefix_query(to_idx);
            else
                return prefix_query(to_idx) - prefix_query(from_idx-1);
        }

        void update(int idx,int add){
            for(++idx; idx < (int)my_arr.size(); idx += (idx & -idx))
                my_arr[idx] += add;              
        }
  
};


void print(BIT const& bit, int length) {
    std::cout << "Index: ";
    for (int idx = 0; idx < length; ++idx)
        std::cout << std::setw(2) << idx << " ";
    std::cout << std::endl;
    std::cout << "Array: ";
    for (int idx = 0; idx < length; ++idx)
        std::cout << std::setw(2) << bit.range_query(idx, idx) << " ";
    std::cout << std::endl << std::endl;

    std::cout << "Prefix sum of first 13 elements: \t" << bit.prefix_query(12) << std::endl;
    std::cout << "Prefix sum of first 7 elements: \t" << bit.prefix_query(6) << std::endl;
    std::cout << "Range sum from index 1 to index 5: \t" << bit.range_query(1, 5) << std::endl;
    std::cout << std::endl;
}

int main()
{
    std::vector<int> array{ 1, 7, 3, 0, 5, 8, 3, 2, 6, 2, 1, 1, 4, 5 };
    int length = array.size();
    BIT bit(array);
    print(bit, length);

    bit.update(4, 2);
    std::cout << "Add 2 to element at index 4" << std::endl << std::endl;
    print(bit, length);
}

https://blog.csdn.net/Yaokai_AssultMaster/article/details/79492190 讲得很清晰

posted on 2020-10-28 12:03  chengyulala  阅读(86)  评论(0编辑  收藏  举报

导航