quick-find.rkt

;; Usage: racket quick-find.rkt < pairs.txt
#lang racket/base

(require racket/list)
(require racket/string)
(require profile)

(provide make-qf-set qf-connected? qf-find qf-count qf-items qf-connect!)
(module+ main (profile-thunk main))

(define (main)
  (let* ([n (string->number (read-line))]
         [set (make-qf-set n)])
    (for ([line (in-lines)])
      (let* ([pair (string-split line " " #:trim? #f)]
             [p (string->number (car pair))]
             [q (string->number (cadr pair))])
        (unless (qf-connected? set p q)
          (qf-connect! set p q)
          (printf "connected ~a to ~a\n" p q))))
    (printf
     "Quick-find Set with ~a connected components: ~a\n"
     (qf-count set)
     (string-join
      (for/list ([n (in-vector (qf-items set))])
        (number->string n))
      #:before-first "[" #:after-last "]"))))

(struct qf-set ([count #:mutable] items) #:transparent)
(define (make-qf-set n) (qf-set n (make-qf-items n)))
(define (make-qf-items n) (list->vector (range n)))

(define (qf-connected? set p q)
  (= (qf-find set p)
     (qf-find set q)))

;; VERY SLOW - it takes over 4 minutes to finish
;; (define (qf-connect! set p q)
;;   (let ([p-root (qf-find set p)]
;;         [q-root (qf-find set q)]
;;         [items (qf-set-items set)])
;;     (unless (= p-root q-root)
;;       (for ([(v i) (in-indexed items)])
;;         (when (= v p-root)
;;           (vector-set! items i q-root)))
;;       (set-qf-set-count! set (- (qf-set-count set) 1)))))

;; better - it takes about 28s to finish
;; but still slower than the Golang version
(define (qf-connect! set p q)
  (let* ([p-root (qf-find set p)]
         [q-root (qf-find set q)]
         [items (qf-set-items set)]
         [n (vector-length items)])
    (unless (= p-root q-root)
      (for ([i (in-range n)] ;; it is weird that adding this line slows down the program by about 10s
            [v (in-vector items)])
        (when (= v p-root) (vector-set! items i q-root)))
      (set-qf-set-count! set (- (qf-set-count set) 1)))))

(define (qf-find set p)
  (let ([items (qf-set-items set)])
    (if (p . < . (vector-length items))
      (vector-ref items p)
      #f)))

(define qf-count qf-set-count)
(define (qf-items set)
  (vector->immutable-vector (qf-set-items set)))