From b3a141753f9331306704818c1b7b806ae8836834 Mon Sep 17 00:00:00 2001
From: Simon Jakobi <simon.jakobi@gmail.com>
Date: Sat, 22 Mar 2014 06:58:39 +0100
Subject: [PATCH] New exercise: python nth-prime

---
 EXERCISES.txt               |  1 +
 nth-prime/example.py        | 31 +++++++++++++++++++++++++++++++
 nth-prime/nth_prime_test.py | 25 +++++++++++++++++++++++++
 3 files changed, 57 insertions(+)
 create mode 100644 nth-prime/example.py
 create mode 100644 nth-prime/nth_prime_test.py

diff --git a/EXERCISES.txt b/EXERCISES.txt
index b00b09fc8d..d7b2f3f76a 100644
--- a/EXERCISES.txt
+++ b/EXERCISES.txt
@@ -21,6 +21,7 @@ gigasecond
 triangle
 scrabble-score
 sieve
+nth-prime
 luhn
 roman-numerals
 binary
diff --git a/nth-prime/example.py b/nth-prime/example.py
new file mode 100644
index 0000000000..d997d8a29b
--- /dev/null
+++ b/nth-prime/example.py
@@ -0,0 +1,31 @@
+from itertools import count
+from math import sqrt
+
+
+def nth_prime(n):
+    known = []
+    candidates = prime_candidates()
+
+    def is_prime(m):
+        sqrt_m = sqrt(m)
+        for k in known:
+            if k > sqrt_m:
+                return True
+            elif m % k == 0:
+                return False
+        return True
+
+    while len(known) < n:
+        x = next(candidates)
+        if is_prime(x):
+            known.append(x)
+
+    return known[n - 1]
+
+
+def prime_candidates():
+    yield 2
+    yield 3
+    for n in count(6, 6):
+        yield n - 1
+        yield n + 1
diff --git a/nth-prime/nth_prime_test.py b/nth-prime/nth_prime_test.py
new file mode 100644
index 0000000000..761e07ce10
--- /dev/null
+++ b/nth-prime/nth_prime_test.py
@@ -0,0 +1,25 @@
+try:
+    from prime import nth_prime
+except ImportError:
+    raise SystemExit('Could not find prime.py. Does it exist?')
+
+import unittest
+
+
+class NthPrimeTests(unittest.TestCase):
+    def test_first_prime(self):
+        self.assertEqual(2, nth_prime(1))
+
+    def test_sixth_prime(self):
+        self.assertEqual(13, nth_prime(6))
+
+    def test_first_twenty_primes(self):
+        self.assertEqual([2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71],
+                         [nth_prime(n) for n in range(1, 21)])
+
+    def test_prime_no_10000(self):
+        self.assertEqual(104729, nth_prime(10000))
+
+
+if __name__ == '__main__':
+    unittest.main()