Segment Trees - Range Maximum Query (RMQ) - Scheme Implementation

segment-tree.rkt

;; Package for Range Maximum Query (RMQ) operations using Segment Trees.
;; Implemented in Racket v5.2.1.
;; Author: Viswanath Sivakumar <[email protected]>

;; Abstraction for integer ranges.
(define make-range
  (lambda (low high) (cons low high)))
(define low
  (lambda (range) (car range)))
(define high
  (lambda (range) (cdr range)))
(define integer-inside-range?
  ;; Checks if an integer falls within a range.
  (lambda (i r) (and (>= i (low r)) (<= i (high r)))))
(define range-inside-range?
  ;; Checks if range r1 falls within range r2.
  (lambda (r1 r2) (and (>= (low r1) (low r2)) (<= (high r1) (high r2)))))

;; Abstraction for segment trees.
;; Each node is represented as a list of value, range and left and right sub-trees.
(define make-segment-tree
  (lambda (value range left right) (list value range left right)))
(define value
  (lambda (seg-tree) (car seg-tree)))
(define range
  (lambda (seg-tree) (cadr seg-tree)))
(define left-child
  (lambda (seg-tree) (caddr seg-tree)))
(define right-child
  (lambda (seg-tree) (cadddr seg-tree)))

;; Create a Range Maximum Query (RMQ) segment tree with all values
;; initialized to 0.
(define build-empty-segment-tree
  ;; Takes an integer n, and returns a segment tree for the range [0, n-1].
  ;; All values are initialized to 0.
  ;; Ranges are 0 based.
  (lambda (n)
    (letrec ((B (lambda (cur-range)
                  (let* ((lo (low cur-range))
                         (hi (high cur-range))
                         (mid (floor (/ (+ lo hi) 2)))
                         (left-range (make-range lo mid))
                         (right-range (make-range (+ mid 1) hi)))
                    (if (= lo hi)
                        (make-segment-tree 0 cur-range '() '())
                        (make-segment-tree 0 cur-range (B left-range) (B right-range)))))))
      (B (make-range 0 (- n 1))))))

;; Create a segment tree for RMQ from a list of integers.
(define build-segment-tree-from-list
  ;; Takes a list of integers and returns a segment tree initialized to the values in the list.
  (lambda (l)
    (letrec ((B (lambda (seg-tree i l)
                  (if (null? l)
                      seg-tree
                      (update-segment-tree (B seg-tree (+ i 1) (cdr l)) i (car l))))))
      (B (build-empty-segment-tree (length l)) 0 l))))

;; Update the ith leaf of the segment tree with a new value
(define update-segment-tree
  ;; Takes as arguments the index i (0-based) and the new value at ith leaf.
  ;; Return the updated segment tree.
  (lambda (seg-tree i new-value)
    (let* ((cur-range (range seg-tree))
           (lo (low cur-range))
           (hi (high cur-range))
           (left (left-child seg-tree))
           (right (right-child seg-tree)))
      (cond ((not (integer-inside-range? i cur-range)) seg-tree)
            ((= lo hi) (make-segment-tree new-value cur-range left right))
            (else (let ((new-left (update-segment-tree left i new-value))
                        (new-right (update-segment-tree right i new-value)))
                    (make-segment-tree (max (value new-left) (value new-right))
                                       cur-range
                                       new-left
                                       new-right)))))))

;; Query the maximum value within a range.
(define query-segment-tree
  ;; Takes a segment tree and a range and returns the maximum value within
  ;; the range.
  (lambda (seg-tree query-range)
    (let* ((cur-range (range seg-tree))
           (lo (low cur-range))
           (hi (high cur-range))
           (lo-query (low query-range))
           (hi-query (high query-range)))
      (cond ((or (> lo-query hi) (< hi-query lo)) #f)
            ((range-inside-range? cur-range query-range) (value seg-tree))
            (else (let ((left-max (query-segment-tree (left-child seg-tree) query-range))
                        (right-max (query-segment-tree (right-child seg-tree) query-range)))
                    (cond ((not (number? left-max)) right-max)
                          ((not (number? right-max)) left-max)
                          (else (max left-max right-max)))))))))

;; Obtain the values in the leaves of the segment tree
(define leaf-values
  ;; Takes a segment tree and returns a list representing
  ;; the values in the leaves of the segment tree.
  (lambda (seg-tree)
    (if (and (null? (left-child seg-tree)) (null? (right-child seg-tree)))
        (cons (value seg-tree) '())
        (append (leaf-values (left-child seg-tree)) (leaf-values (right-child seg-tree))))))

;; Tests
(define (run-tests)
  ;; Build and initialize the tree
  (define seg-tree (build-segment-tree-from-list '(4 7 2 9)))
  (display (leaf-values seg-tree)) (newline) ; (4 7 2 9)
  ;; Queries and updates
  (display (query-segment-tree seg-tree (make-range 0 3))) (newline) ; 9
  (display (query-segment-tree seg-tree (make-range 1 1))) (newline) ; 7
  (display (query-segment-tree seg-tree (make-range 1 2))) (newline) ; 7
  (set! seg-tree (update-segment-tree seg-tree 1 5))
  (display (leaf-values seg-tree)) (newline) ; (4 5 2 9)
  (display (query-segment-tree seg-tree (make-range 1 2))) (newline) ; 5
  (display (query-segment-tree seg-tree (make-range 1 3))) (newline) ; 9
  (set! seg-tree (update-segment-tree seg-tree 3 100))
  (display (leaf-values seg-tree)) (newline) ; (4 5 2 100)
  (display (query-segment-tree seg-tree (make-range 0 2))) (newline) ; 5
  (display (query-segment-tree seg-tree (make-range 0 3))) (newline) ; 100
'done)