Implementation of Dynamic Segment Trees with Poly Hash Tables

// C++ code for the above approach
#include <bits/stdc++.h>
using namespace std;

// This class is used to compute the hash
// of a polynomial with coefficients given
// by the vector a. The hash is
// computed modulo p.
class PolyHash {
public:
    PolyHash(const vector<int>& a, int p)
        : a(a), p(p)
    {
    }

    // Computes the hash of the
    // polynomial at the point k.
    int hash(int k) const
    {
        int result = 0;
        for (int i = 0; i < a.size(); i++) {
            result = (result + a[i] * k) % p;
        }
        return result;
    }

private:
    vector<int> a;

    // Coefficients of the polynomial.
    int p;

    // Modulus used to compute the hash.
};

// This class implements a dynamic
// segment tree with polyhash.
class DynamicSegmentTreeWithPolyHash {
public:
    DynamicSegmentTreeWithPolyHash(const vector<int>& a)
        : a(a)
    {

        // The segment size is chosen to
        // be sqrt(n), where n
        // is the size of a.
        segmentSize = sqrt(a.size());

        // We allocate enough memory for
        // 2*n nodes in the segment tree.
        segmentTree.resize(2 * a.size());

        // We create a vector of
        // unordered_maps, one for
        // each segment of size sqrt(n).
        for (int i = 0; i < a.size(); i += segmentSize) {
            unordered_map<int, int> polyHashTable;
            // We compute the hash of the
            // subarray [i, i+segmentSize]
            // and store it in
            // polyHashTable for each
            // element in that subarray.
            for (int j = i;
                 j < min(i + segmentSize, (int)a.size());
                 j++) {
                polyHashTable[a[j]] = j - i;
                segmentTree[a.size() + i / segmentSize]
                    += a[j];
            }

            // We store the hash table for
            // this segment in polyHashTables.
            polyHashTables.push_back(polyHashTable);
        }
    }

    // Returns the sum of elements in
    // the range [l, r] of the array a.
    int query(int l, int r) const
    {
        int sum = 0;

        // We loop over the elements in
        // the range [l, r] and add
        // them to sum.
        for (int i = l; i <= r;) {

            // If we are at the beginning
            // of a segment, and the segment
            // contains the entire range
            // [i, r], we can add the
            // precomputed sum of the
            // segment to sum and skip
            // to the next segment.
            if (i % segmentSize == 0
                && i + segmentSize - 1 < r) {
                sum += segmentTree[a.size()
                                   + i / segmentSize];
                i += segmentSize;
            }

            // Otherwise, we simply add
            // the current element to sum
            // and move to the next element.
            else {
                sum += a[i];
                i++;
            }
        }
        return sum;
    }

    // Updates the value of a[i] to x.
    void update(int i, int x)
    {

        // We first find the segment
        // that contains a[i].
        int segmentIndex = i / segmentSize;
        int segmentStart = segmentIndex * segmentSize;
        int segmentEnd = min(segmentStart + segmentSize,
                             (int)a.size());

        // We update the value of a[i] in
        // the array and update the sum of
        // the segment in the segment tree.
        segmentTree[a.size() + +segmentIndex] += x - a[i];

        // Update the corresponding node
        // of the segment tree with the
        // difference between the new and
        // old values of the element.
        polyHashTables[segmentIndex].erase(a[i]);
        polyHashTables[segmentIndex][x] = i - segmentStart;
        a[i] = x;
    }

    void insert(int i, int x)

    // In the insert method, a new
    // element x is inserted at index i
    // of the vector a, and then the
    // rebuild method is called to
    // rebuild the segment tree
    // using the updated vector.
    {
        a.insert(a.begin() + i, x);
        rebuild();
    }

    void remove(int i)

    // In the remove method, the element
    // at index i of the vector a is
    // removed, and then the rebuild method
    // is called to rebuild the segment
    // tree using the updated vector.
    {
        a.erase(a.begin() + i);
        rebuild();
    }

private:
    void rebuild()
    {
        int n = a.size();

        // The rebuild method is called in
        // both insert and remove methods
        // to rebuild the segment tree
        // using the updated vector.
        segmentSize = sqrt(n);
        segmentTree.clear();

        // It first calculates the new
        // segmentSize using the
        // updated size of the vector.
        segmentTree.resize(2 * n);
        polyHashTables.clear();

        // Then it clears the segmentTree
        // and polyHashTables vectors and
        // rebuilds them by iterating
        // through the new vector in
        // segments of size segmentSize.
        for (int i = 0; i < n; i += segmentSize) {
            unordered_map<int, int> polyHashTable;
            for (int j = i; j < min(i + segmentSize, n);
                 j++) {
                polyHashTable[a[j]] = j - i;

                // In each segment, it
                // creates a new polyHashTable
                // and adds the values to it,
                // and calculates the sum of
                // the segment and adds
                // it to the segmentTree.
                segmentTree[n + i / segmentSize] += a[j];
            }
            polyHashTables.push_back(polyHashTable);

            // Finally, the polyHashTable is
            // added to the polyHashTables
            // vector.
        }
    }

private:
    vector<int> a;
    int segmentSize;
    vector<int> segmentTree;
    vector<unordered_map<int, int> > polyHashTables;
};

// Driver's code
int main()
{
    vector<int> a = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };

    // Initializes a vector a with 10
    // integers from 1 to 10.
    DynamicSegmentTreeWithPolyHash dst(a);

    // It creates an instance of the
    // DynamicSegmentTreeWithPolyHash
    // class with the vector a.
    cout << "Query(0, 4) = " << dst.query(0, 4) << endl;

    // It calls the query method of the
    // dst object with arguments (0, 4)
    // and prints the result. expected
    // output: 15
    dst.update(3, 100);

    // It calls the update method of the
    // dst object with  arguments (3, 100).
    cout << "Query(2, 5) = " << dst.query(2, 5) << endl;

    // It calls the query method of the
    // dst object with arguments (2, 5)
    // and prints the result. This should
    // output "Query(2, 5) = 114".
    // expected output: 114
    return 0;
}