0834-sum-of-distances-in-tree.rkt

#lang racket

(define/contract (sum-of-distances-in-tree n edges)
  (-> exact-integer? (listof (listof exact-integer?)) (listof exact-integer?))
  (define connect-vec (make-vector n null))
  (define steps-vec (make-vector n 0))
  (define sz-vec (make-vector n 0))
  (define ans-vec (make-vector n 0))
  (for ([e (in-list edges)])
    (define f (first e))
    (define t (second e))
    (vector-set! connect-vec f (cons t (vector-ref connect-vec f)))
    (vector-set! connect-vec t (cons f (vector-ref connect-vec t))))
  (define (fst-steps-sz f t)
    (define-values (steps sz)
      (for/fold ([steps 0] [sz 0] #:result (values steps (add1 sz)))
                ([s (in-list (vector-ref connect-vec t))] #:unless (= s f))
        (define-values (s-steps s-sz) (fst-steps-sz t s))
        (values (+ steps s-steps s-sz) (+ sz s-sz))))
    (vector-set! steps-vec t steps)
    (vector-set! sz-vec t sz)
    (values steps sz))
  (define (snd-steps-sz f t f-steps)
    (define t-steps
      (+ (vector-ref steps-vec t) f-steps (- n (vector-ref sz-vec t))))
    (vector-set! ans-vec t t-steps)
    (for ([s (in-list (vector-ref connect-vec t))] #:unless (= s f))
      (snd-steps-sz t s (- t-steps (vector-ref steps-vec s) (vector-ref sz-vec s)))))
  (fst-steps-sz n 0)
  (snd-steps-sz n 0 0)
  (vector->list ans-vec))

(sum-of-distances-in-tree 6 '((0 1) (0 2) (2 3) (2 4) (2 5)))