Added range query. (#158)

AlfredChester · Mr-Python-in-China · Copilot · web-flow · commit e6ace5c269d6 · 2025-08-26T21:58:59.000+08:00
* Added range query.

* Import range query from __init__.py

* Formatted code according to yapf_style.cfg

* Fixed query io.

* Remove the last endl in to_str

* Fixed range_query_test.py

* write some doc for query.py

* Fix range_query_test.py

* Initialize position_range in cyaron/query.py:get_one_query when nothing is given.

* Set up all supported versions of pythons.

* Use older version of poerty-core

* Fixed expected time complexity.

* Use random.sample to enhance performance. Removed unused parameter in noipstyle.py

* Rollback to while loop to enhance performance.

* Rollback test.yml

* Set TEST_LEN to 20000 to save time.

* Support weight generator in range query.

* Fixed test for weighted range query

* Again fixed tests for weighted range query

* Fixed mismodify to test.yml.

* Added frequency of big queries query.py

* Updated test.yml.

* Fixed a stupid mistake.

* improve type hints

* test python 3.6 on windows

* fix docs

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

---------

Co-authored-by: Mr. Python &lt;2789762371@qq.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -8,26 +8,27 @@ jobs:
   build:
     strategy:
       matrix:
-        platform: [ubuntu-20.04, windows-latest]
+        platform: [ubuntu-22.04, windows-latest]
     runs-on: ${{ matrix.platform }}
     steps:
       - uses: actions/checkout@v3
-      - name: Set up Python 3.6
+      - name: Set up Python 3.6 (Windows)
+        if: matrix.platform == 'windows-latest'
         uses: actions/setup-python@v4
         with:
-          python-version: '3.6'
+          python-version: "3.6"
       - name: Set up Python 3.8
         uses: actions/setup-python@v4
         with:
-          python-version: '3.8'
+          python-version: "3.8"
       - name: Set up Python 3.10
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: "3.10"
       - name: Set up Python 3.12
         uses: actions/setup-python@v4
         with:
-          python-version: '3.12'
+          python-version: "3.12"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
diff --git a/cyaron/__init__.py b/cyaron/__init__.py
@@ -8,7 +8,7 @@
 
 from random import choice, randint, random, randrange, uniform
 
-#from .visual import visualize
+# from .visual import visualize
 from . import log
 from .compare import Compare
 from .consts import *
@@ -21,3 +21,4 @@
 from .string import String
 from .utils import *
 from .vector import Vector
+from .query import RangeQuery
diff --git a/cyaron/graders/noipstyle.py b/cyaron/graders/noipstyle.py
@@ -21,14 +21,18 @@ def noipstyle(content, std):
                                 i + 1, j + 1, content_lines[i][j:j + 5],
                                 std_lines[i][j:j + 5]))
             if len(std_lines[i]) > len(content_lines[i]):
-                return False, TextMismatch(content, std,
-                                           'Too short on line {}.', i + 1,
-                                           j + 1, content_lines[i][j:j + 5],
-                                           std_lines[i][j:j + 5])
+                return False, TextMismatch(
+                    content,
+                    std,
+                    'Too short on line {}.',
+                    i + 1,
+                )
             if len(std_lines[i]) < len(content_lines[i]):
-                return False, TextMismatch(content, std,
-                                           'Too long on line {}.', i + 1,
-                                           j + 1, content_lines[i][j:j + 5],
-                                           std_lines[i][j:j + 5])
+                return False, TextMismatch(
+                    content,
+                    std,
+                    'Too long on line {}.',
+                    i + 1,
+                )
 
     return True, None
diff --git a/cyaron/query.py b/cyaron/query.py
@@ -0,0 +1,238 @@
+"""
+This module provides a `RangeQuery` class for generating queries
+based on limits of each dimension.
+
+Classes:
+    RangeQueryRandomMode: Enum to control how random range endpoints are generated.
+    RangeQuery: A class for generating random queries.
+    
+Usage:
+    n = randint(1, 10)
+    q = randint(1, 10)
+    Q = RangeQuery.random(q, [(1, n)])
+    io.input_writeln(Q)
+"""
+
+import random
+from enum import IntEnum
+from typing import Optional, Union, Tuple, List, Callable, TypeVar, overload, Generic, Any, Sequence
+
+from .utils import list_like
+
+
+class RangeQueryRandomMode(IntEnum):
+    """Control how random range endpoints are generated for range queries."""
+    LESS = 0  # disallow l = r
+    ALLOW_EQUAL = 1  # allow l = r
+
+
+WeightT = TypeVar('WeightT', bound=Tuple[Any, ...])
+
+
+class RangeQuery(Generic[WeightT], Sequence[Tuple[List[int], List[int],
+                                                  WeightT]]):
+    """A class for generating random queries."""
+    result: List[Tuple[List[int], List[int], WeightT]]  # Vector L, R, weights.
+
+    def __init__(self):
+        self.result = []
+
+    def __len__(self):
+        return len(self.result)
+
+    @overload
+    def __getitem__(self, item: int) -> Tuple[List[int], List[int], WeightT]:
+        ...
+
+    @overload
+    def __getitem__(self,
+                    item: slice) -> List[Tuple[List[int], List[int], WeightT]]:
+        ...
+
+    def __getitem__(self, item: Union[int, slice]):
+        return self.result[item]
+
+    def __str__(self):
+        """__str__(self) -> str
+            Return a string to output the queries. 
+            The string contains all the queries with l and r in a row, splits with "\\n".
+        """
+        return self.to_str()
+
+    def to_str(self):
+        """
+        Return a string to output the queries. 
+        The string contains all the queries with l and r (and w if generated) in a row, splits with "\\n".
+        """
+        res = ''
+        for l, r, w in self.result:
+            l_to_str = [str(x) for x in l]
+            r_to_str = [str(x) for x in r]
+            w_to_str = [str(x) for x in w]
+            res += ' '.join(l_to_str) + ' ' + ' '.join(r_to_str)
+            if len(w_to_str) > 0:
+                res += ' ' + ' '.join(w_to_str)
+            res += '\n'
+        return res[:-1]  # remove the last '\n'
+
+    @staticmethod
+    @overload
+    def random(
+        num: int = 1,
+        position_range: Optional[Sequence[Union[int, Tuple[int, int]]]] = None,
+        *,
+        mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+        weight_generator: None = None,
+        big_query: float = 0.2,
+    ) -> "RangeQuery[Tuple[()]]":
+        ...
+
+    @staticmethod
+    @overload
+    def random(
+        num: int = 1,
+        position_range: Optional[Sequence[Union[int, Tuple[int, int]]]] = None,
+        *,
+        mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+        weight_generator: Callable[[int, List[int], List[int]], WeightT],
+        big_query: float = 0.2,
+    ) -> "RangeQuery[WeightT]":
+        ...
+
+    @staticmethod
+    def random(
+        num: int = 1,
+        position_range: Optional[Sequence[Union[int, Tuple[int, int]]]] = None,
+        *,
+        mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+        weight_generator: Optional[Callable[[int, List[int], List[int]],
+                                            WeightT]] = None,
+        big_query: float = 0.2,
+    ):
+        """
+        Generate `num` random queries with dimension limit.
+        Args:
+            num: the number of queries
+            position_range: a list of limits for each dimension
+                single number x represents range [1, x]
+                list [x, y] or tuple (x, y) represents range [x, y]
+            mode: the mode queries generate, see Enum Class RangeQueryRandomMode
+            weight_generator: A function that generates the weights for the queries. It should:
+                - Take the index of query (starting from 1), starting and ending positions as input.
+                - Return a list of weights of any length.
+            big_query: a float number representing the probability for generating big queries.
+        """
+        ret = RangeQuery()
+
+        for i in range(num):
+            ret.result.append(
+                RangeQuery.get_one_query(position_range,
+                                         big_query=big_query,
+                                         mode=mode,
+                                         weight_generator=weight_generator,
+                                         index=i + 1))
+        return ret
+
+    @staticmethod
+    @overload
+    def get_one_query(
+            position_range: Optional[Sequence[Union[int, Tuple[int,
+                                                               int]]]] = None,
+            *,
+            big_query: float = 0.2,
+            mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+            weight_generator: None = None,
+            index: int = 1) -> Tuple[List[int], List[int], Tuple[()]]:
+        ...
+
+    @staticmethod
+    @overload
+    def get_one_query(
+            position_range: Optional[Sequence[Union[int, Tuple[int,
+                                                               int]]]] = None,
+            *,
+            big_query: float = 0.2,
+            mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+            weight_generator: Callable[[int, List[int], List[int]], WeightT],
+            index: int = 1) -> Tuple[List[int], List[int], WeightT]:
+        ...
+
+    @staticmethod
+    def get_one_query(
+            position_range: Optional[Sequence[Union[int, Tuple[int,
+                                                               int]]]] = None,
+            *,
+            big_query: float = 0.2,
+            mode: RangeQueryRandomMode = RangeQueryRandomMode.ALLOW_EQUAL,
+            weight_generator: Optional[Callable[[int, List[int], List[int]],
+                                                WeightT]] = None,
+            index: int = 1):
+        """
+        Generate a pair of query lists (query_l, query_r, w) based on the given position ranges and mode.
+        Args:
+            position_range (Optional[List[Union[int, Tuple[int, int]]]]): A list of position ranges. Each element can be:
+                - An integer, which will be treated as a range from 1 to that integer.
+                - A tuple of two integers, representing the lower and upper bounds of the range.
+            mode (RangeQueryRandomMode): The mode for generating the queries. It can be:
+                - RangeQueryRandomMode.ALLOW_EQUAL: Allow the generated l and r to be equal.
+                - RangeQueryRandomMode.LESS: Ensure that l and r are not equal.
+            weight_generator: A function that generates the weights for the queries. It should:
+                - Take the index of query (starting from 1), starting and ending positions as input.
+                - Return a list of weights of any length.
+        Returns:
+            Tuple[List[int], List[int]]: A tuple containing two lists:
+                - query_l: A list of starting positions.
+                - query_r: A list of ending positions.
+        Raises:
+            ValueError: If the upper-bound is smaller than the lower-bound.
+            ValueError: If the mode is set to less but the upper-bound is equal to the lower-bound.
+        """
+        if position_range is None:
+            position_range = [10]
+
+        dimension = len(position_range)
+        query_l: List[int] = []
+        query_r: List[int] = []
+        for i in range(dimension):
+            cur_range: Tuple[int, int]
+            pr = position_range[i]
+            if isinstance(pr, int):
+                cur_range = (1, pr)
+            elif len(pr) == 1:
+                cur_range = (1, pr[0])
+            else:
+                cur_range = pr
+
+            if cur_range[0] > cur_range[1]:
+                raise ValueError(
+                    "upper-bound should be larger than lower-bound")
+            if mode == RangeQueryRandomMode.LESS and cur_range[0] == cur_range[
+                    1]:
+                raise ValueError(
+                    "mode is set to less but upper-bound is equal to lower-bound"
+                )
+
+            if random.random() < big_query:
+                # Generate a big query
+                cur_l = cur_range[1] - cur_range[0] + 1
+                lb = max(2 if mode == RangeQueryRandomMode.LESS else 1,
+                         cur_l // 2)
+                ql = random.randint(lb, cur_l)
+                l = random.randint(cur_range[0], cur_range[1] - ql + 1)
+                r = l + ql - 1
+            else:
+                l = random.randint(cur_range[0], cur_range[1])
+                r = random.randint(cur_range[0], cur_range[1])
+                # Expected complexity is O(1)
+                # We can use random.sample, But it's actually slower according to benchmarks.
+                while mode == RangeQueryRandomMode.LESS and l == r:
+                    l = random.randint(cur_range[0], cur_range[1])
+                    r = random.randint(cur_range[0], cur_range[1])
+                if l > r:
+                    l, r = r, l
+
+            query_l.append(l)
+            query_r.append(r)
+        if weight_generator is None:
+            return (query_l, query_r, ())
+        return (query_l, query_r, weight_generator(index, query_l, query_r))
diff --git a/cyaron/tests/__init__.py b/cyaron/tests/__init__.py
@@ -5,4 +5,5 @@
 from .compare_test import TestCompare
 from .graph_test import TestGraph
 from .vector_test import TestVector
+from .range_query_test import TestRangeQuery
 from .general_test import TestGeneral
diff --git a/cyaron/tests/range_query_test.py b/cyaron/tests/range_query_test.py
diff --git a/pyproject.toml b/pyproject.toml
