Prime factors in Elixir

In the last article, we showed how pattern matching could solve FizzBuzz. It’s a deliberately simple example so let’s look at something a little bit more complex.

If you recall back to high school math, any non-prime number can be split into the product of two or more smaller primes. So 9 can be split into 3 x 3 and 12 can be split into 2 x 2 x 3.

I like this problem because it lends itself well to TDD. I can write a test for 2 and then for 3 and then for 4 and slowly build the logic up.

To solve the problem for 2, we could write this. There’s nothing really interesting about the code so far.

defmodule PrimeFactors do
  def factor(number) do
    [number]
  end
end

That code also works for 3 so let’s skip to 4.

The way I normally approach factoring 4 is that it’s an even number and by definition any even number can be split into 2 and number/2.

In a language like Java, I might approach it like this.

public static List<Integer> factor(int number) {
  List<Integer> result = new ArrayList<>();
  if( number % 2 == 0 ) { // It's even
    result.add(2);
    number = number / 2;
  }
  result.add(number);
  return result;
}

In elixir, we approach that differently. We declare factor twice, with conditions. If the number is divisible by 2 then we enter the first declaration and if it’s not, we enter the second declaration.

defmodule PrimeFactors do
  def factor(number) when rem(number, 2) == 0 do
    [2, trunc(number/2)]
  end

  def factor(number) do
    [number]
  end
end

The when is called a guard clause and it’s part of Elixir’s pattern matching. If the when clause evaluates to true then it will use that function, otherwise it will keep scanning for the next one that matches.

Now this code works for 4 but fails for 2 because 2 is legitimately divisible by itself. In Java, we might just add another conditional check like this…

public static List<Integer> factor(int number) {
  List<Integer> result = new ArrayList<>();
  if( number % 2 == 0 && number != 2) { // It's even
    result.add(2);
    number = number / 2;
  }
  result.add(number);
  return result;
}

In the elixir code, we’re going to change our approach and start using an accumulator, which is a common pattern.

defmodule PrimeFactors do
  def factor(number) do
    factor(number, [])
  end

  defp factor(number, acc) when number == 2 do
    [number | acc]
  end

  defp factor(number, acc) when rem(number, 2) == 0 do
    number = trunc(number / 2)
    factor(number, [2 | acc])
  end
end

The syntax [number | acc] means that we’re creating a new list with number at the front and the contents of acc coming after it.

Here we’ve introduced a different function with the same name but that takes two parameters instead of one. We’ve also made it private as seen by the defp declaration.

By switching to this accumulator approach, not only did we fix the case of 2 itself, but we’ve solved for every power of two as well. 8, 16, 32, etc.

Of course this is hard coded to powers of 2. What about powers of three or five?

We now need to make the “2” a parameter and also need to be able to increment it to reach all values. For the latter, we need to add one more function variation.

defmodule PrimeFactors do
  def factor(number) do
    factor(number, 2, [])
  end

  defp factor(number, divisor, acc) when number == divisor do
    [number | acc]
  end

  defp factor(number, divisor, acc) when rem(number, divisor) == 0 do
    number = trunc(number / divisor)
    factor(number, divisor, [divisor | acc])
  end

  defp factor(number, divisor, acc) do
    factor(number, divisor + 1, acc)
  end
end

This should now handle every possible positive integer from two up. What about 1? It turns out that 1 will put the code into an infinite loop so we need to add an overload function just to handle that.

defmodule PrimeFactors do
  def factor(1) do
    []
  end
  
  def factor(number) do
    factor(number, 2, [])
  end

  defp factor(number, divisor, acc) when number == divisor do
    [number | acc]
  end

  defp factor(number, divisor, acc) when rem(number, divisor) == 0 do
    number = trunc(number / divisor)
    factor(number, divisor, [divisor | acc])
  end

  defp factor(number, divisor, acc) do
    factor(number, divisor + 1, acc)
  end
end

Now we’re done. Very different from anything we would have written in a non-functional language and yet very clean.

Tags