algorithms/SegmentTree.cpp at master · notepad104/algorithms · GitHub

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
/*
 * C++ Program to Implement Segment Tree
 */
#include <iostream>
#include <cmath>
using namespace std;

// get the middle index from corner indexes.
int getMid(int s, int e)
{
    return s + (e - s) / 2;
}
int getSumUtil(int *st, int ss, int se, int qs, int qe, int index)
{
    if (qs <= ss && qe >= se)
        return st[index];
    if (se < qs || ss > qe)
        return 0;
    int mid = getMid(ss, se);
    return getSumUtil(st, ss, mid, qs, qe, 2 * index + 1) +
           getSumUtil(st, mid+1, se, qs, qe, 2 * index + 2);
}

/* update the nodes which have the given index in
   their range. */
void updateValueUtil(int *st, int ss, int se, int i, int diff, int index)
{
    if (i < ss || i > se)
        return;
    st[index] = st[index] + diff;
    if (se != ss)
    {
        int mid = getMid(ss, se);
        updateValueUtil(st, ss, mid, i, diff, 2 * index + 1);
        updateValueUtil(st, mid + 1, se, i, diff, 2 * index + 2);
    }
}

//update a value in input array and segment tree.
void updateValue(int arr[], int *st, int n, int i, int new_val)
{
    if (i < 0 || i > n-1)
    {
        cout<<"Invalid Input"<<endl;
        return;
    }
    int diff = new_val - arr[i];
    arr[i] = new_val;
    updateValueUtil(st, 0, n-1, i, diff, 0);
}

// Return sum of elements in range from index qs (quey start) to
int getSum(int *st, int n, int qs, int qe)
{
    if (qs < 0 || qe > n-1 || qs > qe)
    {
        cout<<"Invalid Input"<<endl;
        return -1;
    }
    return getSumUtil(st, 0, n-1, qs, qe, 0);
}

// constructs Segment Tree for array[ss..se].
int constructSTUtil(int arr[], int ss, int se, int *st, int si)
{
    if (ss == se)
    {
        st[si] = arr[ss];
        return arr[ss];
    }
    int mid = getMid(ss, se);
    st[si] =  constructSTUtil(arr, ss, mid, st, si * 2 + 1) +
              constructSTUtil(arr, mid + 1, se, st, si * 2 + 2);
    return st[si];
}

/* construct segment tree from given array.  */
int *constructST(int arr[], int n)
{
    int x = (int)(ceil(log2(n)));
    int max_size = 2 * (int)pow(2, x) - 1;
    int *st = new int[max_size];
    constructSTUtil(arr, 0, n-1, st, 0);
    return st;
}
// Main
int main()
{
    int arr[] = {1, 3, 5, 7, 9, 11};
    int n = sizeof(arr)/sizeof(arr[0]);
    int *st = constructST(arr, n);
    cout<<"Sum of values in given range = "<<getSum(st, n, 1, 3)<<endl;
    updateValue(arr, st, n, 1, 10);
    cout<<"Updated sum of values in given range = "<<getSum(st, n, 1, 3)<<endl;
    return 0;
}