disprob.rkt

#lang racket
(require unstable/list)

;; Helper functions

(define (on f g)
  (λ (x y) (f (g x) (g y))))

(define ((<$> f) pair)
    (cons (f (car pair)) (cdr pair)))

;; Tree implementation
(struct tree () #:transparent)

(struct branch tree (node left right) #:transparent)
(struct leaf tree () #:transparent)

(define (get-branch-nodes tree)
  (match tree
    [(leaf) (list (leaf))]
    [(branch _ l r) (list l r)]))

;; Debugging function to show rows of the tree
(define (tree-rows tree)
  (let ([tree-rows (list tree)])
        (letrec ([get-rows
         (λ (cur-row)
           (cond
             [(andmap ((curry equal?) (leaf)) cur-row) '()]
             [else
              (cons (map tree-value cur-row)
                    (get-rows
                     (append* (map get-branch-nodes cur-row))))]
           ))])
    (get-rows tree-rows))))

(define (tree-value tree)
  (match tree
    [(branch n _ _) n]
    [(leaf) (leaf)]))

;; List to Tree stuff

(define (list->tree xs keyfun cmp)
  (let ([sxs (sort xs #:key keyfun cmp)])
    (letrec ([make-tree
    (λ (sxs)
      (match sxs
        [(list x) (branch x (leaf) (leaf))]
        [(list x y) (branch x (leaf) (make-tree (list y)))]
        [(list x y z) (branch y (make-tree (list x)) (make-tree (list z)))]
        [_ (let* ([splitted
             (call-with-values (λ () (split-at sxs
             (floor
             (/ (length sxs) 2)))) list)]
            [mid (last (car splitted))]
            [left (reverse
                   (cdr (reverse (car splitted))))]
            [right (cadr splitted)])
        (branch mid (make-tree left) (make-tree right)))]))])
      (make-tree sxs))))

;; Searching stuff

(define (make-pdist pairs)
  (list->tree pairs car (on > car)))

(define (tree-search tree cmp)
  (cond
    [(leaf? tree) #f]
    [else
     (match (cmp (branch-node tree))
    ['left (tree-search (branch-left tree) cmp)]
    ['right (tree-search (branch-right tree) cmp)]
    [x x])]))

(define (add-dups pairs)
  (map (λ (xs)
       (foldr
        (λ (acc x)
          (cons (+ (car acc) (car x)) (cdr x)))
        (car xs) (cdr xs)))
     (group-by (on equal? cdr) pairs)))

;; Turn weights into probabilities

(define (weights->ps pairs)
  (let ([n (foldr
   (λ (pair acc)
     (+ acc (car pair)))
   (caar pairs)
   (cdr pairs))])
    (map (<$> (λ (p)
                (cond
                  [(= 0 p) 0]
                  [else (/ p n)]))) pairs)))

;; Build the buckets
(define (make-buckets pairs)
  (letrec
      ([gen
      (λ (prev rest)
        (match rest
          [(list (cons p elem))
           (list (cons (cons p 0) elem))]
          [(list-rest (cons p elem) rest)
           (cons (cons (cons prev (- prev p)) elem)
                 (gen (- prev p) rest))]))])
    (gen 1 pairs)))

(define pairs->pdist
  (compose
   make-pdist
   make-buckets
   weights->ps))

(define (between? bucket p)
  (match bucket
    [(cons a b)
     (cond
       [(and
      (< p a)
      (> p b)) #t]
       [else #f])]))

(define ((dist-compare target) tree)
  (match tree
    [(branch (cons mid node) left right)
     (cond
       [(between? mid target) node]
       [(> target (car mid)) 'left]
       [else 'right])]

    [(branch (cons mid node) (leaf) (leaf)) node]

    [(branch (cons mid node) left (leaf))
     (cond
       [(between? mid target) node]
       [else 'left])]
    [(branch (cons mid node) (leaf) right)
     (cond
       [(between? mid target) node]
       [else 'right])]
    [_ 'nil]))

(define (search-pdist cmp pdist)
  (match (cmp pdist)
    ['left (search-pdist cmp (branch-left pdist))]
    ['right (search-pdist cmp (branch-right pdist))]
    [a a]))

;; Runs tests on the probability distribution ``dist'' and returns a hash table
;; with the frequencies of each value
(define (test-pdist n dist)
  (let ([table (make-hasheq)])
  (for ([_ (in-range n)])
    (let ([n (search-pdist (dist-compare (random)) dist)])
    (hash-set! table n (+ 1 (hash-ref table n 0)))))
    table))

(test-pdist 10000 (pairs->pdist (for/list ([i (in-range 20)]) (cons i i))))
(test-pdist 10000 (pairs->pdist (for/list ([i (in-range 11)]) (cons 1 i))))
(test-pdist 10000 (pairs->pdist (for/list ([i (in-range 11)]) (cons 1 i))))
(test-pdist 10000 (pairs->pdist (for/list ([i (in-range 11)]) (cons 1 i))))