rectangle-union-area/segment_tree.cpp

//
// Created by maximilian on 16.04.22.
//
#include <cassert>
#include <bit>
#include <stack>
#include "segment_tree.h"

SegmentTree::SegmentTree(const std::vector<RectCoord> &coords):
   _num_leaf_nodes(coords.size() -1),
    // reserve nodes for a full binary tree with at least as many leaves as intervals
   _num_meta_nodes(std::bit_ceil(_num_leaf_nodes) -1),
   _nodes(2 * _num_meta_nodes + 1, {0, coords[_num_leaf_nodes].coord, coords[_num_leaf_nodes].coord, 0}) {

    assert(!coords.empty());

    // Initialize all nodes with zero coverage and dummy interval

    // We initialize the tree from bottom up, keeping track of the index ranges
    // and the length of the corresponding segment

    for(Index leaf_node_idx = 0 ; leaf_node_idx < _num_leaf_nodes ; ++leaf_node_idx) {
        _nodes[leaf_node_idx + _num_meta_nodes].left_coord = coords[leaf_node_idx].coord;
        _nodes[leaf_node_idx + _num_meta_nodes].right_coord = coords[leaf_node_idx + 1].coord;
    }
    // Note that there are remaining leafs that represent dummy ranges
    for(Index node_idx = _num_meta_nodes - 1; node_idx != -1 ; --node_idx) {
        _nodes[node_idx].left_coord = left_child(node_idx).left_coord;
        _nodes[node_idx].right_coord = right_child(node_idx).right_coord;
    }

    assert(_nodes.front().left_coord == coords.front().coord);
    assert(_nodes.front().right_coord == coords.back().coord);
    assert(_nodes.front().segment_length() == coords.back().coord - coords.front().coord);
}

void SegmentTree::add_interval(Interval interval) {
    check_integrity();
    std::stack<Index> visit_to;
    std::stack<Index> update_to;
    visit_to.push(0);
    while(!visit_to.empty()) {
        Index node_idx = visit_to.top();
        visit_to.pop();
        if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) {
            add_coverage(node_idx);
        } else {
            if(interval.left < left_child(node_idx).right_coord) {
                visit_to.push(left_child_idx(node_idx));
            }
            if(right_child(node_idx).left_coord < interval.right) {
                visit_to.push(right_child_idx(node_idx));
            }
            if(!_nodes[node_idx].covered()) {
                update_to.push(node_idx);
            }
        }
    }
    while(!update_to.empty()) {
        update_covered_length(update_to.top());
        update_to.pop();
    }
    check_integrity();
}

void SegmentTree::remove_interval(Interval interval) {
    check_integrity();
    std::stack<Index> visit_to;
    std::stack<Index> update_to;
    visit_to.push(0);
    while(!visit_to.empty()) {
        Index node_idx = visit_to.top();
        visit_to.pop();
        if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) {
            remove_coverage(node_idx);
        } else {
            if(interval.left < left_child(node_idx).right_coord) {
                visit_to.push(left_child_idx(node_idx));
            }
            if(right_child(node_idx).left_coord < interval.right) {
                visit_to.push(right_child_idx(node_idx));
            }
            if(!_nodes[node_idx].covered()) {
                update_to.push(node_idx);
            }
        }
    }
    while(!update_to.empty()) {
        update_covered_length(update_to.top());
        update_to.pop();
    }
    check_integrity();
}

void SegmentTree::check_integrity() {
    for(Index node_idx = 0 ; node_idx < _nodes.size(); ++node_idx) {
        if(_nodes[node_idx].covered()) {
            assert(_nodes[node_idx].covered_length == _nodes[node_idx].segment_length());
        } else {
            if(is_leaf(node_idx)) {
                assert(_nodes[node_idx].covered_length == 0);
            } else {
                assert(_nodes[node_idx].covered_length == \
                left_child(node_idx).covered_length + right_child(node_idx).covered_length);
            }
        }
    }
}