-
Notifications
You must be signed in to change notification settings - Fork 0
/
math_gen.py
64 lines (52 loc) · 1.81 KB
/
math_gen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import random
import numpy as np
from typing import List
import numexpr as ne
from tqdm import tqdm
ops = ["/", "*", "-", "+"]
nprob = 10000
max_node = 2
pad_len = 10
# calculation node class
class OpNode:
def __init__(self, left: int, right: int, op: str) -> None:
self.left = left
self.right = right
self.op = op
def __str__(self) -> str:
return f"({self.left}{self.op}{self.right})"
def padded_str(self) -> str:
padded_left = "0" * (pad_len - len(str(self.left))) + str(self.left)
padded_right = "0" * (pad_len - len(str(self.right))) + str(self.right)
return f"({padded_left}{self.op}{padded_right})"
def build_tree(num_nodes: int) -> OpNode:
"""
build a problem with num_nodes of random ints
"""
if num_nodes == 1:
num = random.randint(1, 1000)
fac = np.random.choice([1, random.random()], p=[0.6, 0.4])
return num if fac == 1 else round(fac * num, 2)
left_subtree = build_tree(num_nodes // 2)
right_subtree = build_tree(num_nodes - (num_nodes // 2))
op = random.choice(ops)
return OpNode(left_subtree, right_subtree, op)
def generate_problems(nprob: int, max_node: int) -> List[str]:
"""
generate nprob problems with at most max_node ints involved
"""
probs = []
for _ in tqdm(range(nprob)):
num_nodes = random.randint(2, max_node)
prob = build_tree(num_nodes)
try:
ans = str(round(ne.evaluate(str(prob)).item(), 2))
except ZeroDivisionError:
ans = "nan"
# pad prob and ans
ans = "0" * (pad_len - len(ans)) + ans
probs.append(f"${prob.padded_str()}={ans[::-1]}$") # reverse ans
return probs
with open("./test.txt", "w") as f:
for prob in generate_problems(nprob, max_node):
f.write(f"{prob}\n")