Elixirで順に指定個数の素数を列挙する関数をEnum, Stream, Flowで作ってみた

ふと興味を持ったので，エラトステネスのふるいを純朴に利用して，Elixirで小さいものから順に指定個数の素数を列挙する関数をEnum, Stream, Flowを使って作ってみました．

シリーズ

• Elixirで順に指定個数の素数を列挙する関数をNxでも作ってみた

ソースコード全体

Mix.install(
  [
    {:flow, "~> 1.2"},
    {:benchee, "~> 1.1"}
  ]
)

defmodule Prime do

  def prime_candidates() do
    Stream.unfold(2, fn
      2 -> {2, 3}
      n -> {n, n + 2}
    end)
  end

  def prime_enum(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Enum.take(count)
    |> Enum.map(fn {pr, prs} -> {pr, Enum.filter(prs, & rem(pr, &1) == 0)} end)
    |> Enum.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Enum.map(fn {pr, _} -> pr end)
  end

  def prime_stream(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Stream.take(count)
    |> Stream.map(fn {pr, prs} -> {pr, Stream.filter(prs, & rem(pr, &1) == 0)} end)
    |> Stream.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Stream.map(fn {pr, _} -> pr end)
    |> Enum.to_list()
  end

  def prime_flow(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Stream.take(count)
    |> Flow.from_enumerable(max_demand: 1)
    |> Flow.map(fn {pr, prs} -> {pr, Stream.filter(prs, & rem(pr, &1) == 0)} end)
    |> Flow.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Flow.map(fn {pr, _} -> pr end)
    |> Enum.to_list()
  end
end

Benchee.run(
  %{
    "prime_enum" => fn count -> Prime.prime_enum(count) end,
    "prime_stream" => fn count -> Prime.prime_stream(count) end,
    "prime_flow" => fn count -> Prime.prime_flow(count) end
  },
  inputs: %{
    "10" => 10,
    "100" => 100,
    "1000" => 1000,
    "10000" => 10000
  })

コード解説

  def prime_candidates() do
    Stream.unfold(2, fn
      2 -> {2, 3}
      n -> {n, n + 2}
    end)
  end

このプログラム片は次のような列挙をします．

[2, 3, 5, 7, 9, 11, 13, 15, 17, 19, ...]

つまり最初2で，次は3で，以降，次々2を足していって奇数を列挙して，素数の候補(prime_candidates)としています．

次に，

    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)

このプログラム片は，次のような列挙をします．

[
    {2, []},
    {3, [2]},
    {5, [2, 3]},
    {7, [2, 3, 5]},
    ...
]

つまり，素数の候補(prime_candidates)と，それより小さい素数の候補のリストからなるタプルを列挙しています．

次に，

  def prime_enum(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Enum.take(count)

とすることで，引数 count の個数だけ，前述の列挙をします．

Stream版，Flow版だと次のようにします．

  def prime_enum(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Stream.take(count)

次がエラトステネスのふるいの本体となります．

  def prime_enum(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Enum.take(count)
    |> Enum.map(fn {pr, prs} -> {pr, Enum.filter(prs, & rem(pr, &1) == 0)} end)
    |> Enum.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Enum.map(fn {pr, _} -> pr end)
  end

• |> Enum.map(fn {pr, prs} -> {pr, Enum.filter(prs, & rem(pr, &1) == 0)} end)の部分は，素数の候補prを，prより小さい素数の候補のリストprsの各要素で割り切れる数，すなわちprの約数からなるリストを生成して，prとともにタプルにしています．
• |> Enum.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end) で，上記のリストdivisors(約数)の個数が0であるような要素のみからなるリストを生成します．
• |> Enum.map(fn {pr, _} -> pr end)とすることで，前述のようなエラトステネスのふるいを潜り抜けた，素数からなるリストを生成します．

Stream版の場合には次のようにしています．

  def prime_stream(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Stream.take(count)
    |> Stream.map(fn {pr, prs} -> {pr, Stream.filter(prs, & rem(pr, &1) == 0)} end)
    |> Stream.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Stream.map(fn {pr, _} -> pr end)
    |> Enum.to_list()
  end

Flow版の場合には次のようにしています．

  def prime_flow(count) do
    prime_candidates()
    |> Stream.map(fn pr -> {pr, Stream.take_while(prime_candidates(), & &1 < pr)} end)
    |> Stream.take(count)
    |> Flow.from_enumerable(max_demand: 1)
    |> Flow.map(fn {pr, prs} -> {pr, Stream.filter(prs, & rem(pr, &1) == 0)} end)
    |> Flow.filter(fn {_pr, divisors} -> Enum.count(divisors) == 0 end)
    |> Flow.map(fn {pr, _} -> pr end)
    |> Enum.to_list()
  end

ベンチマーク結果(MacStudio on M1 Ultra)

• 10個の時には僅差でEnumが最速でした．
• 100個の時にはStreamが最速でした．
• 1000,10000個の時にはFlowが最速でした．