(ns grep-n)
(use 'clojure.test)
 
(defn grep-n [n f coll]
  (lazy-seq
   (when-let [s (seq coll)]
     (let [p (take-while (complement f) (take n s))
           n-till-first-match (count p)
           matches (take-while f (nthrest s n-till-first-match))]
       (if (seq matches)
         (let [n-till-last-match (+ (count p) (count matches))
               a (take n (nthrest s n-till-last-match))]
           (if (some f a)
             (concat p matches (grep-n n f (nthrest s n-till-last-match)))
             (concat p matches a (grep-n n f (nthrest s (+ n-till-last-match (count a)))))))
         (grep-n n f (rest s)))))))
 
(deftest grep-n-test
  (is (= '(0 1 2)
         (grep-n 1 #{1} (range 10))))
 
  (is (= '(2 3 4)
         (grep-n 1 #{3} (range 10))))
 
  (is (= '(4 5 6 7 8)
         (grep-n 2 #{6} (range 10))))
 
  (is (= '(3 4 5 6 7 12 13 14 15 16)
         (grep-n 2 #{5 14} (range 20))))
 
  (is (= '(0 1 2 3 4 5 6 7 8)
         (grep-n 3 #{3 5} (range 10))))
 
  (is (= '(4 5 6)
         (grep-n 2 #{4} (range 4 10))))
 
  (is (= '(3 4 5 6 7)
         (grep-n 2 #(= 0 (mod % 5)) (range 1 10)))))
 
(run-tests)

(time (doall (take 10 (grep-n 2 #{1 1111 1111111} (range 100000000)))))