        coll))
 
 (defn- counts-flatten
-  "Flatten COLL by K."
+  "Flatten COLL by K.
+
+    (counts-flatten [{:a [{:b 1} {:b 2}], :c 2}
+                     {:a {:b 3}, :c 4}] :a)
+      -> [{:b 1} {:b 2} {:b 3}]"
   [coll k]
   {:pre [(sequential? coll)
          (keyword? k)]}

                 (if (sequential? x)  x
                     [x])))))

(defn- counts-unflatten-1
  "Unflatten a single item in COLL using one count value.
  (Recursively called by `counts-unflatten`)."
  [coll count]
  (condp = count
    nil   [nil (rest coll)]
    :atom [(first coll) (rest coll)]
    :nil  [:nil (rest coll)]
    (split-at count coll)))

(defn- counts-unflatten
  "Unflatten COLL by K using COUNTS from `counts-of`."
  ([coll k counts]
   (counts-unflatten [] coll k counts))
  ([acc coll k [count & more]]
   (let [[unflattend coll] (counts-unflatten-1 coll count)
         acc (conj acc unflattend)]
     (if-not (seq more) (map (fn [x]
                               (when x

;; At a high level, these functions let you aggressively flatten a sequence of maps by a key
;; so you can apply some function across it, and then unflatten that sequence.
;;
;; 1.  Get a value that can be used to unflatten a sequence later with `counts-of`.
;; 2.  Flatten the sequence with `counts-flatten`
;; 3.  Modify the flattened sequence as needed
;; 4.  Unflatten the sequence by calling `counts-unflatten` with the modified sequence value from step 1
;; 5.  `map merge` the original sequence and the unflattened sequence.
;;
;; For your convenience `counts-apply` combines these steps for you.

(defn- counts-of
  "Return a sequence of counts / keywords that can be used to unflatten
   COLL later.

    (counts-of [{:a [{:b 1} {:b 2}], :c 2}
                {:a {:b 3}, :c 4}] :a)
      -> [2 :atom]

   For each `x` in COLL, return:

   *  `(count (k x))` if `(k x)` is sequential
   *  `:atom`         if `(k x)` is otherwise non-nil
