clojure-tree-descendant-counting

(def test-tree {1 {:root nil} 2 {:root 1} 3 {:root 1} 4 {:root 2} 5 {:root 1} 6 {:root 4}})

(defn path-to-top [tree v]
	(if (nil? v)
		'()
		(cons v (path-to-top tree (:root (get tree v))))))

(defn steps-indexed-path
	([upward-path steps]
		(if (= upward-path '())
			'()
			(cons [(first upward-path) steps] (steps-indexed-path (rest upward-path) (+ steps 1)))))
	([upward-path]
		(steps-indexed-path upward-path 0)))

(defn count-descendants [tree]
	(let [markers (reduce concat '() (map steps-indexed-path (map (partial path-to-top tree) (keys tree))))]
		(reduce (fn [counter [vertex generation]] (assoc counter vertex (assoc (get counter vertex {}) generation (+ (get (get counter vertex {}) generation 0) 1)))) {} markers)))

(defn sanitize-descendant-counts [association] 
	(let [max-depth (apply max (keys association))]
		(map (fn [i] (get association i 0)) (range 1 (+ max-depth 1)))))

(defn solve-problem [tree]
	(let [descendant-counts (count-descendants tree)]
		(apply merge (map (fn [v] (hash-map v (vec (sanitize-descendant-counts (get descendant-counts v))))) (keys descendant-counts)))))

(solve-problem test-tree)