Skip to content

Instantly share code, notes, and snippets.

@NolanDeveloper
Last active March 1, 2021 20:36
Show Gist options
  • Save NolanDeveloper/32583beeff39730c5e0152095dab7700 to your computer and use it in GitHub Desktop.
Save NolanDeveloper/32583beeff39730c5e0152095dab7700 to your computer and use it in GitHub Desktop.
Segment tree in pure c
#include <stdlib.h>
#include <stdio.h>
#ifndef PREFIX
#define PREFIX st_
#endif
#ifndef TYPE
#error "Define TYPE for which segment tree should be created. You may also want "
"to define PREFIX for function names which defaults to 'st_'"
#endif
#define _CONCAT(a, b) a ## b
#define CONCAT(a, b) _CONCAT(a, b)
#ifndef ST_FUNCTION
#define ST_FUNCTION(NAME) CONCAT(PREFIX, NAME)
#endif
#ifdef __cplusplus
extern "C" {
#endif
struct SegmentTree {
TYPE *nodes;
size_t array_size;
TYPE (*func)(TYPE, TYPE);
};
/* Create segment tree. func must be associative. */
struct SegmentTree *ST_FUNCTION(create)(TYPE *array, size_t size, TYPE (*func)(TYPE, TYPE));
/* Free memory */
void ST_FUNCTION(delete)(struct SegmentTree * tree);
/* Compute f(a[left], a[left + 1], ..., a[right - 1]) */
TYPE ST_FUNCTION(lookup)(struct SegmentTree *tree, size_t left, size_t right);
/* Update tree value at specific position. */
void ST_FUNCTION(update)(struct SegmentTree *tree, size_t position, TYPE new_value);
#ifdef __cplusplus
}
#endif
#ifdef IMPLEMENTATION
static void ST_FUNCTION(build)(struct SegmentTree *tree, TYPE *array, size_t i, size_t l, size_t r) {
if (l + 1 == r) {
tree->nodes[i] = array[l];
return;
}
size_t m = (l + r) / 2;
ST_FUNCTION(build)(tree, array, 2 * i + 1, l, m);
ST_FUNCTION(build)(tree, array, 2 * i + 2, m, r);
tree->nodes[i] = tree->func(tree->nodes[2 * i + 1], tree->nodes[2 * i + 2]);
}
struct SegmentTree *ST_FUNCTION(create)(TYPE *array, size_t size, TYPE (*func)(TYPE, TYPE)) {
struct SegmentTree *tree = malloc(sizeof(struct SegmentTree));
tree->array_size = size;
tree->nodes = malloc(4 * size * sizeof(TYPE));
tree->func = func;
ST_FUNCTION(build)(tree, array, 0, 0, size);
return tree;
}
void ST_FUNCTION(delete)(struct SegmentTree * tree) {
free(tree->nodes);
free(tree);
}
static TYPE ST_FUNCTION(lookup_helper)(struct SegmentTree *tree, size_t left, size_t right, size_t i, size_t l, size_t r) {
if (left == l && right == r) {
return tree->nodes[i];
}
size_t m = (l + r) / 2;
if (right <= m) {
return ST_FUNCTION(lookup_helper)(tree, left, right, 2 * i + 1, l, m);
} else if (left >= m) {
return ST_FUNCTION(lookup_helper)(tree, left, right, 2 * i + 2, m, r);
}
TYPE a = ST_FUNCTION(lookup_helper)(tree, left, m, 2 * i + 1, l, m);
TYPE b = ST_FUNCTION(lookup_helper)(tree, m, right, 2 * i + 2, m, r);
return tree->func(a, b);
}
TYPE ST_FUNCTION(lookup)(struct SegmentTree *tree, size_t left, size_t right) {
return ST_FUNCTION(lookup_helper)(tree, left, right, 0, 0, tree->array_size);
}
static void ST_FUNCTION(update_helper)(struct SegmentTree *tree, size_t i, size_t l, size_t r, size_t position, TYPE value) {
if (l + 1 == r) {
tree->nodes[i] = value;
return;
}
size_t m = (l + r) / 2;
if (position < m) {
ST_FUNCTION(update_helper)(tree, 2 * i + 1, l, m, position, value);
} else {
ST_FUNCTION(update_helper)(tree, 2 * i + 2, m, r, position, value);
}
}
void ST_FUNCTION(update)(struct SegmentTree *tree, size_t position, TYPE new_value) {
ST_FUNCTION(update_helper)(tree, 0, 0, tree->array_size, position, new_value);
}
#endif
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment