beichen0426/olympiads-ref

1# -----------------------------------------------------------------------------2# Author: Jiawei Liu3# Date: 2025-10-294# -----------------------------------------------------------------------------5import json6import re7from pathlib import Path8from typing import List, Tuple9 10 11problem_tag = "Problem"12solution_tag = "Solution"13# answer_tag = "Answer"14 15 16def segment_exams(text: str):17    matchs = list(18        re.finditer(19            r"^#+\s*.+(?:ASU|CIS)\s*(\d{4})(?:\s*problems)?",20            text,21            flags=re.IGNORECASE | re.MULTILINE,22        )23    )24 25    exams = {}26    for i, m in enumerate(matchs):27        if "cis" in m.group().lower():28            continue29 30        year = m.group(1)31        exam_text = text[32            m.end() : matchs[i + 1].start() if i + 1 < len(matchs) else len(text)33        ]34 35        exams[year] = exam_text.strip()36 37    return exams38 39 40def analyze(text: str) -> Tuple[List, int]:41    """42    Analyze the text and return the tags and problem number.43 44    Args:45        text (str): The markdown text to analyze.46 47    Returns:48        Tuple[List, int]: A tuple containing the tags and problem number.49    """50    problem_pattern = re.compile(r"(?:\n|# )Problem\s+(\d+)", re.IGNORECASE)51    solution_pattern = re.compile(r"(?:\n|# )Solution", re.IGNORECASE)52    # answer_pattern = re.compile(r"(?:\n|# )Answer", re.IGNORECASE)53 54    tags = []55    tags.extend([(x, problem_tag) for x in problem_pattern.finditer(text)])56    problem_num = len(tags)57 58    tags.extend([(x, solution_tag) for x in solution_pattern.finditer(text)])59    # tags.extend([(x, answer_tag) for x in answer_pattern.finditer(text)])60 61    tags.sort(key=lambda x: x[0].start())62    return tags, problem_num63 64 65def segment(text: str, tags):66    starts = []67    ends = []68 69    for i, (m, tag) in enumerate(tags):70        starts.append(tags[i][0].end())71        if i + 1 < len(tags):72            ends.append(tags[i + 1][0].start())73        else:74            ends.append(len(text))75 76    return [77        text[start:end].strip().strip("#").strip() for start, end in zip(starts, ends)78    ]79 80 81def join(tags, segments):82    problem, solution = "", ""83    problem_label, problem_match, solution_match = "", "", ""84    pairs = []85 86    tag_classes = [_[1] for _ in tags]87 88    for (m, tag), (i, segment) in zip(tags, enumerate(segments)):89        if tag == problem_tag:90            problem = segment91            problem_match = m.group(0)92            problem_label = m.group(1)93 94            # Check if there is no solution following this problem95            next_problem_index = 096            try:97                if problem_tag in tag_classes[i + 1 :]:98                    next_problem_index = tag_classes.index(problem_tag, i + 1)99                else:100                    next_problem_index = len(segments)101            except ValueError:102                next_problem_index = len(segments)103 104            if tag_classes[i + 1 : next_problem_index].count(solution_tag) == 0:105                solution = ""106                solution_match = ""107                pairs.append(108                    (problem, solution, problem_label, problem_match, solution_match)109                )110        else:111            solution = segment112            solution_match = m.group(0)113            pairs.append(114                (problem, solution, problem_label, problem_match, solution_match)115            )116 117    return pairs118 119 120def write_pairs(project_root: Path, output_file: Path, pairs):121    output_jsonl_text = ""122    for year, problems in pairs:123        for (124            problem,125            solution,126            problem_label,127            problem_match,128            solution_match,129        ) in problems:130            output_jsonl_text += (131                json.dumps(132                    {133                        "year": year,134                        "tier": "T1",135                        "problem_label": problem_label,136                        "problem_type": None,137                        "exam": "AllSovietUnion",138                        "problem": problem,139                        "solution": solution,140                        "metadata": {141                            "resource_path": output_file.relative_to(142                                project_root143                            ).as_posix(),144                            "problem_match": problem_match,145                            "solution_match": solution_match,146                        },147                    },148                    ensure_ascii=False,149                )150                + "\n"151            )152 153    output_file.write_text(output_jsonl_text, encoding="utf-8")154 155 156if __name__ == "__main__":157    compet_base_path = Path(__file__).resolve().parent.parent158    compet_md_path = compet_base_path / "md"159    seg_output_path = compet_base_path / "segmented"160    project_root = compet_base_path.parent161 162    for md_file in list(compet_md_path.glob("**/*.md")):163        output_file = seg_output_path / md_file.relative_to(compet_md_path).with_suffix(164            ".jsonl"165        )166        output_file.parent.mkdir(parents=True, exist_ok=True)167 168        # Read the markdown file169        markdown_text = md_file.read_text(encoding="utf-8")170 171        # [(year, [(problem, solution, problem_label, problem_match, solution_match), ...]), ...]172        pairs = []173        exams = segment_exams(markdown_text)174        for year, exam_text in exams.items():175            tags, problem_num = analyze(exam_text)176            segments = segment(exam_text, tags)177            inner_pairs = join(tags, segments)178            pairs.append((year, inner_pairs))179 180        write_pairs(project_root, output_file, pairs)181