CodeEval Consecutive Primes

Consider the sequence of natural number, up to a given even element. We have to say in how many ways we can rearrange these numbers so that any couple of neighbors added up gives a prime number. The first element should always be 1, and should go with the last one.
This is the CodeEval problem #187, and here I am giving a Python 3 solution that gets accepted with full marks but does not give any point for the rank, being too slow. Any suggestion for a better algorithm is welcome.

As usual, my first step is converting the given sample in a python test case.
def test_provided_1(self):
    self.assertEqual(1, solution('2'))

def test_provided_2(self):
    self.assertEqual(2, solution('4'))

def test_provided_3(self):
    self.assertEqual(0, solution('5'))

def test_sixteen(self):
    self.assertEqual(81024, solution('16'))

def test_eighteen(self):
    self.assertEqual(770144, solution('18'))
Do not even try the last one until you have a good implementation at hand, because it is going to take a long time to give you back an answer.
Notice that the third test violates the requirement stating that we should expect an even number in input. I interpretate it as a suggestion to explicitly take care of these cases, avoiding the main computation, as better described by the code below:
def solution(line):
    size = int(line)
    if size % 2:
        return 0
    return necklaces([x for x in range(1, size + 1)], 1)
If the passed size is odd, I just return zero. In any case, if you think about it, an odd number of numbers in the natural sequence can't be reorganized as requested by the problem.
I build on the fly, by list comprehension, the sequence I need, starting from 1 ending to size.
I called the function necklaces because, in the metaphor used in the problem, we want to count how many necklaces we can generate from the passed pebbles.

A brute force solution is absolutely out of scope, since generating all the possible pebbles permutations is a factorial algorithm. You could try it, using the itertools permutations function, but keep your input very low.

To save some CPU time, I decided to cache the prime numbers that my script should be aware of in a set:
PRIMES = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31}
Actually, this won't be a substantial help. Even for a simple prime checker function, the number involved are not big enough to cause performance problems.

After some thinking I came up with this recursive function:
def necklaces(beads, pos):  # 1
    if len(beads) == pos:  # 2
        return 1 if 1 + beads[- 1] in PRIMES else 0

    result = 0

    if beads[pos] + beads[pos-1] in PRIMES:  # 3
        result += necklaces(beads, pos + 1)

    for i in range(pos+2, len(beads), 2):  # 4
        if beads[pos-1] + beads[i] in PRIMES:  # 5
            beads[i], beads[pos] = beads[pos], beads[i]
            result += necklaces(beads, pos + 1)
            beads[i], beads[pos] = beads[pos], beads[i]

    return result
1. I pass to the function the list of beads, and the position we have to check. The first time I call it, from the solution() function, beads is the natural sequence from 1 to the number passed in input, and pos is 1, since we have as a requisite that the first one in the list, position 0, won't change.
2. We have reached the last pebble. Check the sum of its value, added to 1 for the first in the list. If this is a prime number, return 1.
3. Check the sum of the current pebble with the previous one. If it is a prime, accept it and call recursively the function on the next pebble.
4. We need to check all the other pebbles in this position. Actually, not really all of them, since we know that we are looking for a prime number. So, if the previous pebble was even, now we want an odd one, and viceversa. This means, we can safely skip by two in the loop.
5. If the current 'i' pebble makes a good pair with the previous one, we accept it and check for the other ones, just like as above in (3). With the variation that we have to put the 'i' pebble in the 'pos' position, and moving somewhere else the 'pos' pebble. Best place I could think of, was the one resulting from a swap between the two pebbles. This has the advantage that I can easily swap them back after the recursive call is done. An alternative would be using two different lists. One for the actual necklace, the other for the currently unused pebbles.

I couldn't think of a better algorithm, still, the only way to get marks from it on CodeEval was reimplementing it in C++. It was easy, just a few minor changes, and the reported time costs improved by a factor x20.

Even if a bit disappointed, I pushed test case and python script to GitHub.

No comments:

Post a Comment