CODEWITHSUNDEEP
clojure
Marcus Junius Brutus
Marcus Junius Brutus

Reputation: 27286

Clojure take-while and n more items

What's the idiomatic way in Clojure to implement take-while-and-n-more below:

=> (take-while-and-n-more #(<= % 3)  1 (range 10))
(0 1 2 3 4)

My try is:

(defn take-while-and-n-more [pred n coll]
  (let
      [take-while-result (take-while pred coll)
       n0 (count take-while-result)] 
    (concat
     take-while-result
     (into [] (take n (drop n0 coll))))))

Upvotes: 7

Views: 1250

Answers (3)

Leon Grapenthin
Leon Grapenthin

Reputation: 9266

The following code is a modified version of Clojures take-while. Where Clojures take-while returns nil as a default case (when the predicate does not match), this one invokes take to take the the additional items after the predicate fails.

Note that unlike versions using split-with, this version traverses the sequence only once.

(defn take-while-and-n-more
  [pred n coll]
  (lazy-seq
   (when-let [s (seq coll)]
     (if (pred (first s))
       (cons (first s) (take-while-and-n-more pred n (rest s)))
       (take n s)))))

Upvotes: 2

leonardoborges
leonardoborges

Reputation: 5619

Yet another way:

(defn take-while-and-n-more [pred n coll]
  (let [[a b] (split-with pred coll)]
    (concat a (take n b))))

Upvotes: 3

soulcheck
soulcheck

Reputation: 36767

I would use split-with, which is equivalent of getting results of both take-while and drop-while for the same parameters:

(defn take-while-and-n-more [pred n coll]
    (let [[head tail] (split-with pred coll)]
         (concat head (take n tail))))

Upvotes: 10

Related Questions