Search code examples
clojure

Is there a better way to write this pivot function in Clojure?


I've written a function to "pivot" some data tables but wonder if there's a simpler way to accomplish the same result.

(defn pivot-ab-ba
  [data]
  (r/reduce
   (fn [result [ks c]]
     (assoc-in result ks c))
   {}
   ;; pivot {a [b1 c1 b2 c2]} => [[b1 a] c1] [[b2 a] c2]
   (mapcat (fn [[a bcseq]]
             ;; pivot [a [b c]] => [[[b a] c]]
             (mapcat (fn [[b c]] [[[b a] c]]) bcseq))
           data)))

(let [data {1 {:good [1 2] :bad [3 4]}
            2 {:good [5 6] :bad [7 8]}}]
  (pivot-ab-ba data))

; => {:good {1 [1 2], 2 [5 6]}, :bad {1 [3 4], 2 [7 8]}}

This works, but seems over-complicated.

UPDATE:

@TaylorWood proposed a solution below; here is that answer with a modification to avoid passing in the keys that are being pivoted:

(defn pivot [data]
  (reduce-kv
    (fn [acc k v]
      (reduce (fn [acc' k'] (assoc-in acc' [k' k] (k' v)))
              acc
              (keys v)))
    {}
    data))

UPDATE 2: Thank you all for your answers. Because there is such a diversity of answers I profiled the results to get an idea about how they performed. Admittedly, this is a single test, but still interesting:

Benchmarks performed with (criterium.core/bench pivot-function)

# Original pivot-ab-ba
Evaluation count : 8466240 in 60 samples of 141104 calls.
Execution time mean : 7.274613 µs
Execution time std-deviation : 108.681498 ns

# @TaylorWood - pivot
Evaluation count : 39848280 in 60 samples of 664138 calls.
Execution time mean : 1.568971 µs
Execution time std-deviation : 32.567822 ns

# @AlanThompson - reorder-tree
Evaluation count : 25999260 in 60 samples of 433321 calls.
Execution time mean : 2.385929 µs
Execution time std-deviation : 33.130731 ns

# @AlanThompson reorder-tree-reduce
Evaluation count : 14507820 in 60 samples of 241797 calls.
Execution time mean : 4.249135 µs
Execution time std-deviation : 89.933197 ns

# @amalloy - pivot
Evaluation count : 12721980 in 60 samples of 212033 calls.
Execution time mean : 5.087314 µs
Execution time std-deviation : 226.242206 ns

Solution

  • Here's another way to do it:

    (defn pivot [data ks]
      (reduce-kv
        (fn [acc k v]
          (reduce (fn [acc' k'] (assoc-in acc' [k' k] (k' v)))
                  acc
                  ks))
        {}
        data))
    

    This takes a map and the expected keys, then reduces over each key/value pair in data, then does another inner reduce over each expected key, grabbing their value from the data map and assoc'ing it into the output map.

    (def data
      {1 {:good [1 2] :bad [3 4]}
       2 {:good [5 6] :bad [7 8]}})
    
    user=> (pivot data [:good :bad])
    {:good {1 [1 2], 2 [5 6]}, :bad {1 [3 4], 2 [7 8]}}