Merge branch 'main' into rich/ab

reasoning_gym/algorithmic/__init__.py

@@ -10,6 +10,7 @@ from .ab import ABConfig, ABDataset
 from .base_conversion import BaseConversionConfig, BaseConversionDataset
 from .binary_matrix import BinaryMatrixConfig, BinaryMatrixDataset
 from .caesar_cipher import CaesarCipherConfig, CaesarCipherDataset
+from .count_primes import CountPrimesConfig, CountPrimesDataset
 from .group_anagrams import GroupAnagramsConfig, GroupAnagramsDataset
 from .isomorphic_strings import IsomorphicStringsConfig, IsomorphicStringsDataset
 from .letter_counting import LetterCountingConfig, LetterCountingDataset
@@ -69,4 +70,6 @@ __all__ = [
     "BinaryMatrixDataset",
     "ABConfig",
     "ABDataset",
+    "CountPrimesConfig",
+    "CountPrimesDataset",
 ]

reasoning_gym/algorithmic/count_primes.py

@@ -0,0 +1,63 @@
+"""Count prime numbers in a given interval.
+
+Solution obtained with Sieve of Eratosthenes:
+https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes
+"""
+
+import math
+from dataclasses import dataclass
+from random import Random
+from typing import Optional
+
+from ..factory import ProceduralDataset, register_dataset
+
+QUESTION_TEMPLATE = """Count how many prime numbers there are between {start} and {end} (inclusive) ?"""
+
+
+@dataclass
+class CountPrimesConfig:
+    """Configuration for Count Primes dataset generation"""
+
+    max_n: int = 10_000  # Upper bound for the interval
+
+    size: int = 500  # Virtual dataset size
+    seed: Optional[int] = None
+
+    def validate(self):
+        """Validate configuration parameters"""
+        assert 1 <= self.max_n, "max_n must be at least 1"
+
+
+class CountPrimesDataset(ProceduralDataset):
+    """Generates Count Primes exercises with configurable difficulty"""
+
+    def __init__(self, config: CountPrimesConfig):
+        super().__init__(config=config, seed=config.seed, size=config.size)
+        self.primes = self._get_primes(config.max_n + 1)
+
+    def _get_primes(self, n: int) -> list[bool]:
+        if n <= 1:
+            return []
+        primes = [True] * n
+        primes[0] = primes[1] = False
+        for i in range(2, int(math.sqrt(n)) + 1):
+            if primes[i]:
+                for j in range(2 * i, n, i):
+                    primes[j] = False
+        return primes
+
+    def __getitem__(self, idx: int) -> dict:
+        """Generate a single Count Primes question"""
+        rng = Random(self.seed + idx)
+        start = rng.randint(1, self.config.max_n)
+        end = rng.randint(start, self.config.max_n)
+        primes = self.primes[start : end + 1]
+        answer = sum(primes)
+        return {
+            "question": QUESTION_TEMPLATE.format(start=start, end=end),
+            "answer": str(answer),
+            "metadata": {"start": start, "end": end, "primes": primes, "solution": answer},
+        }
+
+
+register_dataset("count_primes", CountPrimesDataset, CountPrimesConfig)

reasoning_gym/algorithmic/rotate_matrix.py

@@ -60,22 +60,16 @@ class RotateMatrixDataset(ProceduralDataset):
         matrix = [numbers[i * n : (i + 1) * n] for i in range(n)]
         return matrix
 
+    def _rot90(self, matrix: list[list[int]]) -> list[list[int]]:
+        """quarter clockwise rotation"""
+        return [list(row) for row in zip(*matrix[::-1])]
+
     def _get_rotated(self, matrix: list[list[int]], num_rotations: int) -> list[list[int]]:
         """Rotate the matrix K times by 90 degrees clockwise"""
         num_rotations %= 4
-        n = len(matrix)
         output = deepcopy(matrix)
-
         for _ in range(num_rotations):
-            for l in range(n // 2):
-                for i in range(l, n - 1 - l):
-                    (output[l][i], output[i][n - 1 - l], output[n - 1 - l][n - 1 - i], output[n - 1 - i][l]) = (
-                        output[n - 1 - i][l],
-                        output[l][i],
-                        output[i][n - 1 - l],
-                        output[n - 1 - l][n - 1 - i],
-                    )
-
+            output = self._rot90(output)
         return output
 
     def _matrix_to_str(self, matrix: list[list[int]]) -> str: