Insert sum of all parent branches in each branch of a nested tree structure

Question

In a tree structure, how can a key/value be merged with each branch where the value is the sum of the branch's value and all the parent branch values?

For example, starting with follow tree structure:

[{  :v 1
    :_ [{ :v 2 }
        { :v 3
          :_ [{ :v 5 }]}
        { :v 4 }]}]

how could it be recreated as:

[{  :v 1
    :sum 1
    :_ [{ :v 2
          :sum 3 }
        { :v 3
          :sum 4
          :_ [{ :v 5
                :sum 9 }]}
        { :v 4
          :sum 5 }]}]

I have been trying with walk. I think this could be the correct approach. But, so far I have not managed it.

Answer 1

I think this recursive function does the trick.

(defn sums [{v :v children :_} sum]
  {:v v
   :sum (+ sum v)
   :_ (mapv #(sums % (+ sum v)) children)})

When evaluated in the following way:

(def root
   [{:v 1
      :_ [{:v 2}
          {:v 3
           :_ [{:v 5}]}
          {:v 4}]}])    

[(sums (first root) 0)]

Results in:

[{:v 1,
  :sum 1,
  :_ [{:v 2,
       :sum 3,
       :_ []}
      {:v 3,
       :sum 4,
       :_ [{:v 5,
            :sum 9,
            :_ []}]}
      {:v 4,
       :sum 5,
       :_ []}]}]

Or here's a more friendly version of the same sums function for your tree format.

(defn sums [root]
  (letfn [(f [{v :v children :_} sum]
            {:v v
             :sum (+ sum v)
             :_ (mapv #(f % (+ sum v)) children)})]
    [(f (first root) 0)]))

(sums root)
;= same result as before

Answer 2

For the sake of comparison, here's a version that uses clojure.walk. I think this is one situation where a custom-built recursive function is going to be cleaner than using walk. A custom function allows you to pass intermediate results (in the form of the sum of the parents) from parent down to child as a function parameter, whereas the function that you pass to walk has no extra parameters other than the form being walked, so you have to record intermediate results in the data itself as you traverse the tree.

(require '[clojure.walk :as walk])

(defn sums
  [x]
  (walk/prewalk (fn [m]
                  (if (map? m)
                    (let [v (or (:v m) 0)
                          s (+ v (or (:sum m) 0))
                          m (assoc m :sum s)]
                      (if (seq (:_ m))
                        (update-in m [:_]
                                   (partial map (fn [c] (assoc c :sum s))))
                        m))
                    m))
                x))

Answer 3

A zipper solution could be written

(require '[clojure.zip :as z])

(def root
  [{:v 1
    :_ [{:v 2}
        {:v 3
         :_ [{:v 5}]}
        {:v 4}]}])    

(def zipper (partial z/zipper :_ :_ (fn [n ch] (assoc n :_ (vec ch)))))

(loop [node (zipper (first root))] 
  (if (z/end? node) 
    (z/root node) 
    (recur 
      (z/next 
       (z/edit node 
                 #(assoc % :sum ((fnil + 0) %2 (:v %))) 
                 (some-> node z/up z/node :sum))))))

;=> {:v 1, 
     :_ [{:v 2, :sum 3}
         {:v 3, 
          :_ [{:v 5, :sum 9}], 
          :sum 4} 
         {:v 4, :sum 5}], 
     :sum 1}

which you could wrap back in a vector if desired. Note, "root" as defined does not match the rest of the tree structure when wrapped in a vector, which is why we have a "(first root)".