线段树源码

线段树源码

---

SegmentTree.h源码

#pragma once
#include "stdafx.h"
class SegmentTreeNode
{
public:
    int min;
    int max;
    double coverWeight;
    double containWeight; //部分覆盖+全部覆盖
    SegmentTreeNode* left;
    SegmentTreeNode* right;
    SegmentTreeNode(int min, int max):left(NULL),right(NULL), coverWeight(0), containWeight(0)
    {this->min = min; this->max = max;}
};
class SegmentTree
{
private:
    SegmentTreeNode* root;
private:
    void add(SegmentTreeNode* node, int min, int max, double weight)
    {
        if(node->min == min && node->max == max) //全部覆盖了，就停止二分
        {
            node->coverWeight += weight; //在getWeight函数里处理部分覆盖时，需要用全部覆盖的权重
            node->containWeight += (max-min)*weight; //加上全部覆盖的权重
            return;
        }
        node->containWeight += (max-min)*weight; //部分覆盖，加上部分覆盖的权重和

        if(node->min +1 >= node->max) 
        {
            return;
        }
        int mid = (node->min+node->max)/2;
        if(node->left == NULL)
        {
            node->left = new SegmentTreeNode(node->min,mid);
        }
        if(node->right == NULL)
        {
            node->right = new SegmentTreeNode(mid, node->max);
        }
        if( max <= mid) //线段在mid的左边
        {
            add(node->left, min,max,weight);
        }
        else if(min >= mid)//线段在mid的右边
        {
            add(node->right, min, max, weight);
        }
        else //mid将线段一分为二
        {
            add(node->left, min, mid, weight);
            add(node->right, mid, max, weight);
        }
    }

    double getWeight(SegmentTreeNode* node, int min, int max)//复杂度O(N),N为树的深度
    {
        if(node == NULL)
        {
            return 0;
        }
        if(node->min == min && node->max == max) //完全覆盖，containWeight就表示在这个区间上的所有权重和
        {
            return node->containWeight;
        }
        double w = node->coverWeight*(max-min);
        int mid = (node->min+node->max)/2;
        if(max <= mid)//线段在mid的左边
        {

            w += getWeight(node->left,min, max); 
        }
        else if( min >= mid) //线段在mid的右边
        {
            w += getWeight(node->right,min,max);
        }
        else //mid将线段一分为二
        {
            w += getWeight(node->left,min, mid) + getWeight(node->right,mid,max);
        }
        return w;
    }
    void destroy(SegmentTreeNode* node)
    {
        if(node == NULL)
        {
            return;
        }
        destroy(node->left);
        destroy(node->right);
        delete node;
    }
public:
    SegmentTree(int min, int max)
    {
        root = new SegmentTreeNode(min,max);
    }
    ~SegmentTree()
    {
        if(root!= NULL)
        {
            destroy(root);
            root = NULL;
        }
    }

    void add(int min, int max, double weight=1)
    {
        add(root, min,max,weight);
    }
    double getWeight(int min, int max)
    {
        return getWeight(root, min, max);
    }

};

SegmentTreeTest.h源码

#pragma once
#include "stdafx.h"
#include "SegmentTree.h"
#include <iostream>
using namespace std;
class SegmentTreeTest
{
public:
    void doTest()
    {
        SegmentTree* segmentTree = new SegmentTree(-10,10);
        segmentTree->add(-3,5,1);
        segmentTree->add(-5,10,3);
        segmentTree->add(3,7,5);
        double weight = segmentTree->getWeight(-1,5);
        cout << weight;
    }
};