diff --git a/reasoning_gym/algorithmic/cryptarithm.py b/reasoning_gym/algorithmic/cryptarithm.py new file mode 100644 index 00000000..5331c83d --- /dev/null +++ b/reasoning_gym/algorithmic/cryptarithm.py @@ -0,0 +1,201 @@ +""" +Cryptarithm puzzle generator (numbers -> letters approach). + +Generates puzzles such that: + WORD1 + + WORD2 + [+ WORD3] + --------- + RESULT +where each letter corresponds to exactly one digit (0..9). +No leading letter can be zero (unless allow_leading_zero=True). +""" + +from dataclasses import dataclass +from random import Random +from typing import Optional + +from ..factory import ProceduralDataset, register_dataset + +EXAMPLE_CASE = """ + BASE ++ BALL +------ + GAMES + +Answer (one possible solution): + +B=7, A=8, S=2, E=9, L=1, G=1, M=0 +Summation: 7829 + 7811 = 15640 (the puzzle might produce a different arrangement, but the principle is the same).""" + +@dataclass +class CryptarithmConfig: + """Configuration for Cryptarithm dataset generation.""" + min_words: int = 2 # Minimum number of addends + max_words: int = 3 # Maximum number of addends + allow_leading_zero: bool = False + seed: Optional[int] = None + size: int = 20 # Number of puzzle instances to generate + include_example: bool = False + + def validate(self): + """Validate configuration parameters.""" + assert 2 <= self.min_words <= self.max_words, \ + "min_words must be <= max_words, both >= 2." + assert self.size > 0, "Dataset size must be positive." + +class CryptarithmDataset(ProceduralDataset): + """ + Generates cryptarithm puzzles by: + 1) Randomly choosing integers for each "addend" (with no leading zero if not allowed), + 2) Summing them, + 3) Mapping distinct digits (0..9) to letters (A..Z), + 4) Formatting the puzzle text. + + This approach guarantees sum correctness and avoids repeated failures. + """ + def __init__(self, config: CryptarithmConfig): + super().__init__(config=config, seed=config.seed, size=config.size) + + def __getitem__(self, idx: int) -> dict: + rng = Random(self.seed + idx) + return self._create_single_puzzle(rng) + + def _create_single_puzzle(self, rng: Random) -> dict: + """ + Creates one puzzle with N addends (2..3) plus a result. + Ensures total distinct digits <= 10. + """ + # 1) Pick how many addends + n_words = rng.randint(self.config.min_words, self.config.max_words) + + # 2) For each addend, pick a random length (3..5) and then pick a random integer with that many digits. + # If leading zero is disallowed, the first digit is from 1..9. + word_lengths = [rng.randint(3, 5) for _ in range(n_words)] + words_numbers = [] + for length in word_lengths: + if self.config.allow_leading_zero: + # e.g. random integer in [0, 10^length - 1], then zero-pad to length + num = rng.randint(0, 10**length - 1) + else: + # leading digit is from 1..9, rest are from 0..9 + # e.g. random integer in [10^(length-1), 10^length - 1] + num = rng.randint(10**(length - 1), 10**length - 1) + words_numbers.append(num) + + # 3) Compute the sum + total_sum = sum(words_numbers) + # The sum can have up to (max_length+1) digits, which is normal in cryptarithms. + + # 4) Gather all digits from the addends and the sum + digits_in_use = set() + def collect_digits(num: int): + return set(str(num)) + + for wn in words_numbers: + digits_in_use.update(collect_digits(wn)) + digits_in_use.update(collect_digits(total_sum)) + + # If we exceed 10 distinct digits, try again (pick new random numbers). + # In practice, we can loop until success. But for demonstration, let's do a simple re-pick approach. + # We'll do a while loop up to some attempts: + if len(digits_in_use) > 10: + # Just do a recursion call to pick new numbers, ignoring current picks + return self._create_single_puzzle(rng) + + # 5) Map each digit to a letter + # If no leading zero is allowed, the leading digit of each addend + result must not map to '0'. + # Actually, we are generating real numeric values, so there's no scenario of leading "0" for + # the addends we enforced (except if allow_leading_zero is True). + # For the puzzle's perspective, we simply create a random assignment from {digits_in_use} -> letters. + # Then the solver has to figure it out. They don't see the digits, only letters. + + digits_in_use_list = sorted(list(digits_in_use)) # e.g. ['0', '1', '3', '9'] + rng.shuffle(digits_in_use_list) # shuffle so mapping is random + letters_pool = [chr(i) for i in range(ord('A'), ord('Z') + 1)] + rng.shuffle(letters_pool) + chosen_letters = letters_pool[: len(digits_in_use_list)] + + # digit -> letter mapping + digit_to_letter = {} + for d, letter in zip(digits_in_use_list, chosen_letters): + digit_to_letter[d] = letter + + # If leading-zero is not allowed, we must ensure that the first digit of each addend and the sum + # does not map to the letter that is assigned to digit '0'. If we see a conflict, we can just re-pick + # or we can try to swap letters. The simplest is to re-pick for demonstration. + if not self.config.allow_leading_zero and '0' in digit_to_letter: + zero_letter = digit_to_letter['0'] + # Check the first digit of each addend and of the sum + for wn in words_numbers: + first_digit = str(wn)[0] + if digit_to_letter.get(first_digit) == zero_letter: + # Conflict => re-generate puzzle + return self._create_single_puzzle(rng) + sum_first_digit = str(total_sum)[0] + if digit_to_letter.get(sum_first_digit) == zero_letter: + return self._create_single_puzzle(rng) + + # Now we have a stable digit->letter mapping. Let's create the letter->digit mapping for the answer. + letter_to_digit = {v: int(k) for k, v in digit_to_letter.items()} + + # 6) Convert each integer to its letter representation + def int_to_letter_str(num: int) -> str: + return "".join(digit_to_letter[d] for d in str(num)) + + words_letters = [int_to_letter_str(num) for num in words_numbers] + result_letters = int_to_letter_str(total_sum) + + # 7) Create the puzzle text + # We'll do the typical vertical format, with a plus sign before the last addend, dashes, then result + puzzle_lines = [] + max_width = max(len(w) for w in words_letters + [result_letters]) + for i, wl in enumerate(words_letters): + if i < len(words_letters) - 1: + # Right align with spaces, +2 for the " " prefix + puzzle_lines.append(f"{wl:>{max_width+2}}") + else: + # Right align with spaces, +2 for the "+ " prefix + puzzle_lines.append(f"+ {wl:>{max_width}}") + + # The line of dashes should match the longest line + puzzle_lines.append("-" * (max_width + 2)) + # Right align the result + puzzle_lines.append(f"{result_letters:>{max_width+2}}") + + puzzle_text = "\n".join(puzzle_lines) + + question_str = ( + "Solve this cryptarithm:\n\n" + f"{puzzle_text}\n\n" + "Each letter stands for a unique digit (0-9). " + + ( + "Leading letters may be zero.\n" + if self.config.allow_leading_zero + else "No leading letter can be zero.\n" + ) + + "Provide a mapping from letters to digits that satisfies the equation.\n" + ) + if self.config.include_example: + question_str += "Here's an example:\n"+EXAMPLE_CASE + + # 8) Create a human-readable answer, e.g. "A=1,B=0,C=9,..." + sorted_letter_keys = sorted(letter_to_digit.keys()) + answer_str = ",".join(f"{letter}={letter_to_digit[letter]}" for letter in sorted_letter_keys) + + # 9) Return the final puzzle item + return { + "question": question_str, + "answer": answer_str, + "metadata": { + "letters": list(letter_to_digit.keys()), + "word_values": words_numbers, + "sum_number": total_sum, + "words_letters": words_letters, + "result_letters": result_letters, + "digit_to_letter": digit_to_letter, + "letter_to_digit": letter_to_digit, + }, + } + +register_dataset("cryptarithm", CryptarithmDataset, CryptarithmConfig) diff --git a/tests/test_cryptarithm.py b/tests/test_cryptarithm.py new file mode 100644 index 00000000..c18fb1e7 --- /dev/null +++ b/tests/test_cryptarithm.py @@ -0,0 +1,98 @@ +import pytest +from reasoning_gym import create_dataset +from reasoning_gym.algorithmic.cryptarithm import CryptarithmDataset, CryptarithmConfig + +def test_cryptarithm_generation(): + dataset = create_dataset("cryptarithm", seed=42, size=10) + assert isinstance(dataset, CryptarithmDataset) + unique_number = set() + for item in dataset: + # Check required keys exist + assert "question" in item + assert "answer" in item + assert "metadata" in item + + # Validate question format + question = item["question"] + assert "Solve this cryptarithm:" in question + assert "Each letter stands for a unique digit (0-9)" in question + + # Validate metadata structure + metadata = item["metadata"] + assert "letters" in metadata + assert "letter_to_digit" in metadata + assert "words_letters" in metadata + assert "result_letters" in metadata + assert "word_values" in metadata + assert "sum_number" in metadata + + # Validate letter to digit mapping + letter_to_digit = metadata["letter_to_digit"] + used_digits = set(letter_to_digit.values()) + assert len(used_digits) == len(letter_to_digit), "Each letter should map to a unique digit" + assert all(0 <= digit <= 9 for digit in used_digits), "All digits should be between 0 and 9" + + # Validate the arithmetic + word_values = metadata["word_values"] + result_value = metadata["sum_number"] + assert sum(word_values) == result_value, "Sum of word values should equal result value" + unique_number.add(result_value) + + assert len(unique_number) == len(dataset) + +def test_cryptarithm_config(): + # Test invalid configs raise assertions + with pytest.raises(AssertionError): + dataset = create_dataset("cryptarithm", min_words=1) # min_words must be >= 2 + + with pytest.raises(AssertionError): + dataset = create_dataset("cryptarithm", min_words=4, max_words=3) # min must be <= max + + with pytest.raises(AssertionError): + dataset = create_dataset("cryptarithm", size=0) # size must be positive + +def test_leading_zero_constraint(): + # Test with leading zeros not allowed + dataset = create_dataset("cryptarithm", seed=42, size=5, allow_leading_zero=False, max_words=10, min_words=5) + + for item in dataset: + # print(item['question']) + metadata = item["metadata"] + letter_to_digit = metadata["letter_to_digit"] + words_letters = metadata["words_letters"] + result_letters = metadata["result_letters"] + + # Check leading letters of all words and result + leading_letters = [word[0] for word in words_letters] + [result_letters[0]] + for letter in leading_letters: + assert letter_to_digit[letter] != 0, "Leading letters cannot be zero when allow_leading_zero=False" + +def test_deterministic_generation(): + dataset1 = create_dataset("cryptarithm", seed=42, size=5) + dataset2 = create_dataset("cryptarithm", seed=42, size=5) + + for i in range(5): + assert dataset1[i]["question"] == dataset2[i]["question"] + assert dataset1[i]["answer"] == dataset2[i]["answer"] + assert dataset1[i]["metadata"] == dataset2[i]["metadata"] + +def test_word_length_constraints(): + dataset = create_dataset("cryptarithm", seed=42, size=10) + + for item in dataset: + metadata = item["metadata"] + words_letters = metadata["words_letters"] + + # Check each word is between 3-5 letters as specified in the code + for word in words_letters: + assert 3 <= len(word) <= 5, "Each word should be between 3 and 5 letters long" + +def test_max_letters_constraint(): + dataset = create_dataset("cryptarithm", seed=42, size=10) + + for item in dataset: + metadata = item["metadata"] + letter_to_digit = metadata["letter_to_digit"] + + # Check total unique letters doesn't exceed 10 (digits 0-9) + assert len(letter_to_digit) <= 10, "Total unique letters should not exceed 10" \ No newline at end of file