mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
143 lines
4.6 KiB
Python
143 lines
4.6 KiB
Python
"""Print elements of a matrix in spiral order.
|
|
|
|
A popular Leetcode problem:
|
|
https://leetcode.com/problems/spiral-matrix/description/
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from random import Random
|
|
from typing import Any, Optional
|
|
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
|
|
QUESTION_TEMPLATE = """Given a matrix, your job is to generate a list of elements in spiral order, starting from the top-left element.
|
|
|
|
Example:
|
|
- Input: For the matrix below, what is the list of elements in spiral order?
|
|
1 2 3
|
|
4 5 6
|
|
7 8 9
|
|
- Output: 1 2 3 6 9 8 7 4 5
|
|
- Explanation:
|
|
- We start from the top-left element (1) and move right until we reach the end of the row: 1 2 3
|
|
- Then, we move down until we reach the last column: 1 2 3 6 9
|
|
- Next, we move left until we reach the first column: 1 2 3 6 9 8 7
|
|
- Then, we move up until we reach the second row (i.e. one below the previously traversed row): 1 2 3 6 9 8 7 4
|
|
- Finally, we move right until we reach the second to last column: 1 2 3 6 9 8 7 4 5
|
|
- The output format is a space-separated list of elements in spiral order (as opposed to a python list)
|
|
|
|
For the matrix below, what is the list of elements in spiral order?
|
|
{matrix}
|
|
"""
|
|
|
|
|
|
@dataclass
|
|
class SpiralMatrixConfig:
|
|
"""Configuration for Spiral Matrix dataset generation"""
|
|
|
|
max_n: int = 10 # Maximum number of rows/cols in the matrix
|
|
|
|
size: int = 500 # Virtual dataset size
|
|
seed: Optional[int] = None
|
|
|
|
def validate(self):
|
|
"""Validate configuration parameters"""
|
|
assert 2 <= self.max_n, "max_n must be at least 2"
|
|
|
|
|
|
class SpiralMatrixDataset(ProceduralDataset):
|
|
"""Generates Spiral Matrix exercises with configurable difficulty"""
|
|
|
|
def __init__(self, config: SpiralMatrixConfig):
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
def _get_matrix(self, rng: Random) -> list[list[int]]:
|
|
"""Generate a random matrix"""
|
|
n = rng.randint(2, self.config.max_n)
|
|
numbers = [rng.randint(0, 9) for _ in range(n**2)]
|
|
rng.shuffle(numbers)
|
|
matrix = [numbers[i * n : (i + 1) * n] for i in range(n)]
|
|
return matrix
|
|
|
|
def _get_spiral(self, matrix: list[list[int]]) -> list[int]:
|
|
"""Return the elements of the matrix in spiral order"""
|
|
t, b = 0, len(matrix)
|
|
l, r = 0, len(matrix[0])
|
|
|
|
out = []
|
|
|
|
while True:
|
|
for i in range(l, r):
|
|
out.append(matrix[t][i])
|
|
t += 1
|
|
if t == b:
|
|
break
|
|
|
|
for i in range(t, b):
|
|
out.append(matrix[i][r - 1])
|
|
r -= 1
|
|
if l == r:
|
|
break
|
|
|
|
for i in range(r - 1, l - 1, -1):
|
|
out.append(matrix[b - 1][i])
|
|
b -= 1
|
|
if t == b:
|
|
break
|
|
|
|
for i in range(b - 1, t - 1, -1):
|
|
out.append(matrix[i][l])
|
|
l += 1
|
|
if l == r:
|
|
break
|
|
|
|
return out
|
|
|
|
def _matrix_to_str(self, matrix: list[list[int]]) -> str:
|
|
"""Get a string representation of the matrix"""
|
|
return "\n".join(" ".join(str(x) for x in row) for row in matrix)
|
|
|
|
def _list_to_str(self, array: list[int]) -> str:
|
|
"""Get a string representation of the array"""
|
|
return " ".join(str(x) for x in array)
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
"""Generate a single Spiral Matrix question"""
|
|
rng = Random(self.seed + idx)
|
|
|
|
matrix = self._get_matrix(rng)
|
|
matrix_str = self._matrix_to_str(matrix)
|
|
answer = self._get_spiral(matrix)
|
|
answer_str = self._list_to_str(answer)
|
|
|
|
return {
|
|
"question": QUESTION_TEMPLATE.format(matrix=matrix_str),
|
|
"answer": answer_str,
|
|
"metadata": {"matrix": matrix, "solution": answer},
|
|
}
|
|
|
|
def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
|
|
"""Overwrite this method in derived classes if a single oracle answer is not available."""
|
|
oracle_answer = entry["answer"].strip()
|
|
|
|
if answer is not None and len(answer) > 0:
|
|
answer = answer.strip()
|
|
|
|
# Exact match
|
|
if answer == oracle_answer:
|
|
return 1.0
|
|
|
|
# Try to see if the model's answer is a python list
|
|
try:
|
|
answer = " ".join(str(item) for item in eval(answer))
|
|
if answer == oracle_answer:
|
|
return 0.5
|
|
else:
|
|
return 0.01
|
|
except Exception as e:
|
|
return 0.01
|
|
|
|
return 0.0
|
|
|
|
|
|
register_dataset("spiral_matrix", SpiralMatrixDataset, SpiralMatrixConfig)
|