open-thought/reasoning-gym
1
0
Fork 0
mirror of https://github.com/open-thought/reasoning-gym.git synced 2026-04-19 12:58:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add new probability problems dataset and extend combinatorics with additional task types

Browse source
This commit is contained in:
Ritvik19 2026-04-18 19:26:10 +05:30
parent 9847d71dce
commit dc0d81c096
5 changed files with 978 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

251
reasoning_gym/combinatorics/combinatorics.py
View file

@ -1,5 +1,6 @@
import math
import random
from collections import Counter
from dataclasses import dataclass, field
from typing import Any, Optional
@ -8,7 +9,24 @@ from ..factory import ProceduralDataset, register_dataset
DATASET_NAME = "combinatorics"
TASK_TYPES = ("ncr", "npr", "permutations_repetition", "inclusion_exclusion", "stars_and_bars", "pigeonhole")
TASK_TYPES = (
"ncr",
"npr",
"permutations_repetition",
"inclusion_exclusion",
"stars_and_bars",
"pigeonhole",
"multinomial",
"grid_paths",
"constrained_selection",
"circular_permutation",
"geometric_counting",
"dictionary_rank",
"derangement",
"group_division",
"legendres_formula",
"integral_solutions",
)
@dataclass
@ -16,7 +34,13 @@ class CombinatoricsConfig:
min_n: int = 5
max_n: int = 15
task_types: tuple[str, ...] = TASK_TYPES
task_weights: list[float] = field(default_factory=lambda: [0.2, 0.15, 0.2, 0.2, 0.15, 0.1])
task_weights: list[float] = field(
default_factory=lambda: [
0.08, 0.06, 0.08, 0.08, 0.06, 0.04,
0.07, 0.07, 0.07, 0.07, 0.07, 0.07,
0.06, 0.06, 0.06, 0.06,
]
)
seed: Optional[int] = None
size: int = 500
@ -119,6 +143,219 @@ class CombinatoricsDataset(ProceduralDataset):
)
return {"question": question, "answer": str(answer), "task_type": "pigeonhole"}
# --- Advanced Counting Principles ---
def _make_multinomial(self, rng: random.Random) -> dict:
num_vars = rng.randint(2, 4)
n = rng.randint(self.config.min_n, self.config.max_n)
parts = self._random_partition(rng, n, num_vars)
var_names = ["x", "y", "z", "w"][:num_vars]
numerator = math.factorial(n)
denominator = 1
for p in parts:
denominator *= math.factorial(p)
answer = numerator // denominator
term_strs = [f"{v}^{e}" for v, e in zip(var_names, parts)]
sum_str = " + ".join(var_names)
question = (
f"What is the coefficient of {' * '.join(term_strs)} in the expansion of "
f"({sum_str})^{n}? Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "multinomial"}
@staticmethod
def _random_partition(rng: random.Random, n: int, k: int) -> list[int]:
"""Generate a random composition of n into k positive parts."""
if k == 1:
return [n]
cuts = sorted(rng.sample(range(1, n), k - 1))
parts = [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [n - cuts[-1]]
return parts
def _make_grid_paths(self, rng: random.Random) -> dict:
m = rng.randint(2, self.config.max_n)
n = rng.randint(2, self.config.max_n)
answer = math.comb(m + n, m)
question = (
f"How many shortest paths are there from the top-left corner to the bottom-right corner "
f"of a {m} x {n} grid, if you can only move right or down? "
f"Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "grid_paths"}
def _make_constrained_selection(self, rng: random.Random) -> dict:
total_men = rng.randint(3, max(4, self.config.max_n))
total_women = rng.randint(3, max(4, self.config.max_n))
committee_size = rng.randint(3, min(total_men + total_women - 1, 8))
min_women = rng.randint(1, min(total_women, committee_size - 1))
answer = 0
for w in range(min_women, min(total_women, committee_size) + 1):
men_needed = committee_size - w
if men_needed > total_men:
continue
answer += math.comb(total_women, w) * math.comb(total_men, men_needed)
question = (
f"A committee of {committee_size} people is to be formed from {total_men} men and "
f"{total_women} women. If at least {min_women} woman/women must be included, how many "
f"ways can the committee be formed? Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "constrained_selection"}
# --- Special Permutations & Geometry ---
def _make_circular_permutation(self, rng: random.Random) -> dict:
n = rng.randint(self.config.min_n, self.config.max_n)
identical_rotations = rng.choice([True, False])
if identical_rotations:
answer = math.factorial(n - 1) // 2
question = (
f"How many distinct ways can {n} people be seated around a circular table, "
f"where clockwise and counter-clockwise arrangements are considered the same? "
f"Give your answer as a single integer."
)
else:
answer = math.factorial(n - 1)
question = (
f"How many distinct ways can {n} people be seated around a circular table? "
f"Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "circular_permutation"}
def _make_geometric_counting(self, rng: random.Random) -> dict:
sub_type = rng.choice(["triangles", "diagonals"])
if sub_type == "triangles":
n = rng.randint(max(6, self.config.min_n), max(7, self.config.max_n))
m = rng.randint(3, n - 3)
answer = math.comb(n, 3) - math.comb(m, 3)
question = (
f"There are {n} points in a plane, of which {m} are collinear. "
f"How many distinct triangles can be formed using these points as vertices? "
f"Give your answer as a single integer."
)
else:
n = rng.randint(max(4, self.config.min_n), max(5, self.config.max_n))
answer = n * (n - 3) // 2
question = (
f"How many diagonals does a {n}-sided convex polygon have? "
f"Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "geometric_counting"}
def _make_dictionary_rank(self, rng: random.Random) -> dict:
length = rng.randint(3, min(6, max(4, self.config.max_n)))
letters = sorted(rng.sample("ABCDEFGHIJKLMNOPQRSTUVWXYZ", length))
word_letters = letters[:]
rng.shuffle(word_letters)
word = "".join(word_letters)
rank = 1
remaining = sorted(word_letters)
for i, ch in enumerate(word):
pos = remaining.index(ch)
rank += pos * math.factorial(len(remaining) - 1)
remaining.pop(pos)
question = (
f"If all permutations of the letters {', '.join(sorted(set(word)))} are arranged "
f"in alphabetical (dictionary) order, what is the rank (position) of the word '{word}'? "
f"Give your answer as a single integer."
)
return {"question": question, "answer": str(rank), "task_type": "dictionary_rank"}
# --- Distribution & Partitioning ---
def _make_derangement(self, rng: random.Random) -> dict:
n = rng.randint(self.config.min_n, min(self.config.max_n, 12))
answer = self._subfactorial(n)
question = (
f"How many derangements (permutations where no element appears in its original position) "
f"are there of a set of {n} elements? Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "derangement"}
@staticmethod
def _subfactorial(n: int) -> int:
if n == 0:
return 1
if n == 1:
return 0
d_prev2, d_prev1 = 1, 0
for i in range(2, n + 1):
d_curr = (i - 1) * (d_prev1 + d_prev2)
d_prev2, d_prev1 = d_prev1, d_curr
return d_prev1
def _make_group_division(self, rng: random.Random) -> dict:
num_groups = rng.randint(2, 4)
n = rng.randint(max(self.config.min_n, num_groups * 2), max(self.config.min_n + 1, self.config.max_n))
group_sizes = self._random_partition(rng, n, num_groups)
group_sizes.sort(reverse=True)
numerator = math.factorial(n)
denominator = 1
for g in group_sizes:
denominator *= math.factorial(g)
size_counts = Counter(group_sizes)
for cnt in size_counts.values():
if cnt > 1:
denominator *= math.factorial(cnt)
answer = numerator // denominator
sizes_str = ", ".join(str(s) for s in group_sizes)
question = (
f"In how many ways can {n} distinct items be divided into unlabeled groups of sizes "
f"{sizes_str}? Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "group_division"}
# --- Number Theory in Combinatorics ---
def _make_legendres_formula(self, rng: random.Random) -> dict:
n = rng.randint(self.config.min_n, self.config.max_n)
primes = [p for p in [2, 3, 5, 7, 11, 13] if p <= n]
if not primes:
primes = [2]
p = rng.choice(primes)
exponent = 0
pk = p
while pk <= n:
exponent += n // pk
pk *= p
question = (
f"What is the largest power of {p} that divides {n}!? "
f"In other words, find the largest k such that {p}^k divides {n}!. "
f"Give your answer as a single integer (the value of k)."
)
return {"question": question, "answer": str(exponent), "task_type": "legendres_formula"}
def _make_integral_solutions(self, rng: random.Random) -> dict:
r = rng.randint(2, 5)
variant = rng.choice(["non_negative", "positive"])
n = rng.randint(max(self.config.min_n, r), self.config.max_n)
if variant == "non_negative":
answer = math.comb(n + r - 1, r - 1)
var_list = " + ".join(f"x{i+1}" for i in range(r))
question = (
f"How many non-negative integer solutions are there to the equation "
f"{var_list} = {n}? Give your answer as a single integer."
)
else:
answer = math.comb(n - 1, r - 1)
var_list = " + ".join(f"x{i+1}" for i in range(r))
question = (
f"How many positive integer solutions are there to the equation "
f"{var_list} = {n}? Give your answer as a single integer."
)
return {"question": question, "answer": str(answer), "task_type": "integral_solutions"}
def __getitem__(self, idx: int) -> dict:
rng = random.Random(self.seed + idx)
task_type = rng.choices(self.config.task_types, weights=self.config.task_weights, k=1)[0]
@ -130,6 +367,16 @@ class CombinatoricsDataset(ProceduralDataset):
"inclusion_exclusion": self._make_inclusion_exclusion,
"stars_and_bars": self._make_stars_and_bars,
"pigeonhole": self._make_pigeonhole,
"multinomial": self._make_multinomial,
"grid_paths": self._make_grid_paths,
"constrained_selection": self._make_constrained_selection,
"circular_permutation": self._make_circular_permutation,
"geometric_counting": self._make_geometric_counting,
"dictionary_rank": self._make_dictionary_rank,
"derangement": self._make_derangement,
"group_division": self._make_group_division,
"legendres_formula": self._make_legendres_formula,
"integral_solutions": self._make_integral_solutions,
}
result = generators[task_type](rng)
return {