Instantly share code, notes, and snippets.

Embed
What would you like to do?
Lazy list implementation for Mathematica
(* Project Euler Problem 1 *)
(* Find the sum of all the multiples of 3 or 5 below 1000. *)
(* StreamSource[#&] is a simple idiom for "a stream of all the natural numbers". *)
ints = TakeWhile[StreamSource[# &], # < 1000 &];
(* Choose only the numbers that are multiples of 3 or 5. *)
filtered = Select[ints, Mod[#, 3] == 0 || Mod[#, 5] == 0];
(* Sum them. Will probably add Stream/:Total[] since it comes up a lot. *)
Fold[Plus, 0, filtered]
(* Project Euler Problem 2 *)
(* By considering the terms in the Fibonacci sequence whose values do not exceed four million, find the sum of the even-valued terms. *)
(* Create a stream of Fibonacci numbers. *)
fibo = StreamSource[Fibonacci];
(* Use TakeWhile to get all the numbers under 4 million, and then filter for evenness. *)
filtered = fibo ~TakeWhile~ (#<4000000&) ~Select~ EvenQ;
(* Total them up. *)
Fold[Plus, 0, filtered]
(* Project Euler Problem 10 *)
(* Find the sum of all the primes below two million. *)
(* Create a stream of Prime numbers. *)
primes = StreamSource[Prime];
(* Sum the ones under 2 million. *)
Fold[Plus, 0, Select[primes, #<2000000&]]
(*
This project arose from my attempt to solve some of the Project Euler
problems using Mathematica. I found it frustrating that I had to guess
how many prime numbers or how many Fibonacci numbers I'd have to scan
in order to satisfy the "primes under 2 million" or "Fibonacci numbers
under 4 million".
I knew how I'd solve the problem in C# or Python (chaining a bunch of
IEnumerables or generators, respectively), and wanted to see how close
I could get to that style in Mathematica, so I posed the question on
the fledgling Mathematica Stack Exchange site.
A helpful user named WReach wrote up a great rough implementation[1],
citing his experience with the Haskell programming language as his
major source of inspiration.
I cleaned up his implementation, put it in a package, and made it look
a little more native (First, Rest instead of Head, Tail). I've added a
few more functions (TakeWhile, FoldList, Most, Last).
Finally, I added a helper called StreamSource that lets you trivially
turn many built-ins (Fibonacci[] and Prime[], in fact) into lazy
sources that you can pump for as many items as you need without going
to too much effort.
I'm providing this code under an MIT-style license in the hopes that
others find it useful. If anyone ends up adding any new features,
especially by providing Stream upvalues for more built-ins, please
send me a pull request.
Happy hacking!
[1]: http://mathematica.stackexchange.com/a/885/178
*)
(*
Copyright (c) 2012 Scott Blomquist.
Permission is hereby granted, free of charge, to any person
obtaining a copy of this software and associated documentation
files (the "Software"), to deal in the Software without
restriction, including without limitation the rights to use, copy,
modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*)
BeginPackage["LazyList`"];
LazyList::usage = "A wrapper type that holds a lazy list."
EmptyQ::usage = "A predicate that determines if a given LazyList is empty."
LazySource::usage = "A lazy list generator that takes a source such as Prime[] or Fibonacci[]."
Begin["`Private`"];
Unprotect[LazyList,EmptyQ]
SetAttributes[LazyList, {HoldAll}]
EmptyQ[LazyList[]] := True
EmptyQ[LazyList[_,_]] := False
LazyList/:First[LazyList[h_,_]] := h
LazyList/:Rest[LazyList[_,t_]] := t
LazyList/:Most[LazyList[h_,t_]] := If[EmptyQ[t],t,LazyList[h,Most[t]]]
LazyList/:Last[z_LazyList] := First[NestWhile[Rest, z, !EmptyQ[Rest[#]]&]]
LazyList/:Part[z_LazyList,0] := LazyList
LazyList/:Part[z_LazyList,1] := First[z]
LazyList/:Part[z_LazyList,n_Integer] := Part[Rest[z],n-1]
LazyList/:Take[z_LazyList, 0] := LazyList[]
LazyList/:Take[z_LazyList, n_] /; n > 0 :=
With[{nn = n-1}, LazyList[First[z], Take[Rest[z], nn]]]
LazyList/:TakeWhile[z_LazyList,crit_] := If[!crit[First[z]],LazyList[],LazyList[First[z],TakeWhile[Rest[z],crit]]]
LazyList/:List[z_LazyList] :=
Module[{tag}
, Reap[
NestWhile[(Sow[First[#], tag]; Rest[#])&, z, !EmptyQ[#]&]
, tag
][[2]] /. {l_} :> l
]
LazyList/:Map[_, LazyList[]] := LazyList[]
LazyList/:Map[fn_, z_LazyList] := LazyList[fn[First[z]], Map[fn, Rest[z]]]
LazyList/:Select[z_LazyList, pred_] :=
NestWhile[Rest, z, (!EmptyQ[#] && !pred[First[#]])&] /.
LazyList[h_, t_] :> LazyList[h, Select[t, pred]]
LazyList/:Fold[fun_,x0_,z0_LazyList] :=
NestWhile[{fun[#[[1]],First[#[[2]]]],Rest[#[[2]]]}&,{x0,z0},!EmptyQ[#[[2]]]&][[1]]
LazyList/:FoldList[_,_,LazyList[]] := LazyList[]
LazyList/:FoldList[fun_,x0_,z0_LazyList] :=
fun[x0,First[z0]]/.x1_:>LazyList[x1,FoldList[fun,x1,Rest[z0]]]
LazyList[{}] := LazyList[]
LazyList[lst_List] := LazyList[Evaluate[First[lst]],LazyList[Evaluate[Rest[lst]]]]
LazyList[f_] := LazySource[f,1]
LazySource[f_, n_:1] := With[{nn = n + 1}, LazyList[f[n], LazySource[f, nn]]]
Protect[LazyList,EmptyQ]
End[]
EndPackage[]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment