random_tree

ソースコード

from titan_pylib.algorithm.random.random_tree import RandomTreeType
from titan_pylib.algorithm.random.random_tree import RandomTree

view on github

展開済みコード

  1# from titan_pylib.algorithm.random.random_tree import RandomTree
  2# from titan_pylib.data_structures.heap.deletable_min_heap import DeletableMinHeap
  3# from titan_pylib.data_structures.heap.min_heap import MinHeap
  4# from titan_pylib.my_class.supports_less_than import SupportsLessThan
  5from typing import Protocol
  6
  7
  8class SupportsLessThan(Protocol):
  9
 10    def __lt__(self, other) -> bool: ...
 11from typing import TypeVar, Generic, Iterable
 12
 13T = TypeVar("T", bound=SupportsLessThan)
 14
 15
 16class MinHeap(Generic[T]):
 17
 18    def __init__(self, a: Iterable[T] = []) -> None:
 19        self.a = list(a)
 20        self._heapify()
 21
 22    def _heapify(self) -> None:
 23        for i in range(len(self.a) - 1, -1, -1):
 24            self._down(i)
 25
 26    def _down(self, i: int) -> None:
 27        a = self.a
 28        n = len(a)
 29        while i * 2 + 1 < n:
 30            u, v = i * 2 + 1, i * 2 + 2
 31            if v < n and a[u] > a[v]:
 32                u = v
 33            if a[i] > a[u]:
 34                a[i], a[u] = a[u], a[i]
 35                i = u
 36            else:
 37                break
 38
 39    def _up(self, i: int) -> None:
 40        a = self.a
 41        while i > 0:
 42            p = (i - 1) >> 1
 43            if a[i] < a[p]:
 44                a[i], a[p] = a[p], a[i]
 45                i = p
 46            else:
 47                break
 48
 49    def get_min(self) -> T:
 50        return self.a[0]
 51
 52    def pop_min(self) -> T:
 53        res = self.a[0]
 54        self.a[0] = self.a[-1]
 55        self.a.pop()
 56        self._down(0)
 57        return res
 58
 59    def push(self, key: T) -> None:
 60        self.a.append(key)
 61        self._up(len(self.a) - 1)
 62
 63    def pushpop_min(self, key: T) -> T:
 64        if self.a[0] > key or self.a[0] == key:
 65            return key
 66        res = self.a[0]
 67        self.a[0] = key
 68        self._down(0)
 69        return res
 70
 71    def replace_min(self, key: T) -> T:
 72        res = self.a[0]
 73        self.a[0] = key
 74        self._down(0)
 75        return res
 76
 77    def __getitem__(self, k: int) -> T:
 78        assert k == 0
 79        return self.a[0]
 80
 81    def tolist(self) -> list[T]:
 82        return sorted(self.a)
 83
 84    def __len__(self):
 85        return len(self.a)
 86
 87    def __str__(self):
 88        return str(self.a)
 89
 90    def __repr__(self):
 91        return f"{self.__class__.__name__}({self})"
 92# from titan_pylib.my_class.supports_less_than import SupportsLessThan
 93from typing import TypeVar, Generic, Iterable
 94
 95T = TypeVar("T", bound=SupportsLessThan)
 96
 97
 98class DeletableMinHeap(Generic[T]):
 99
100    def __init__(self, a: Iterable[T] = []) -> None:
101        """削除可能Minヒープです。
102        要素の 追加/削除/最小値取得 が効率よく行えます。
103        """
104        self.hq: MinHeap[T] = MinHeap(a)
105        self.rem_hq: MinHeap[T] = MinHeap()
106        self._len: int = len(self.hq)
107
108    def push(self, key: T) -> None:
109        """``key`` を追加します。
110        :math:`O(\\log{n})` です。
111        """
112        self._len += 1
113        if self.rem_hq and self.rem_hq.get_min() == key:
114            self.rem_hq.pop_min()
115            return
116        self.hq.push(key)
117
118    def remove(self, key: T) -> None:
119        """``key`` を削除します。
120        :math:`O(\\log{n})` です。
121        """
122        assert self._len > 0
123        self._len -= 1
124        if self.hq.get_min() == key:
125            self.hq.pop_min()
126        else:
127            self.rem_hq.push(key)
128
129    def _delete(self) -> None:
130        while self.rem_hq and self.rem_hq.get_min() == self.hq.get_min():
131            self.hq.pop_min()
132            self.rem_hq.pop_min()
133
134    def get_min(self) -> T:
135        """最小の要素を返します。
136        :math:`O(\\log{n})` です。
137        """
138        assert self._len > 0
139        self._delete()
140        return self.hq.get_min()
141
142    def pop_min(self) -> T:
143        """最小の要素を削除して返します。
144        :math:`O(\\log{n})` です。
145        """
146        assert self._len > 0
147        self._len -= 1
148        self._delete()
149        return self.hq.pop_min()
150
151    def __len__(self) -> int:
152        return self._len
153import enum
154from typing import Optional
155import random
156
157
158class RandomTreeType(enum.Enum):
159    """``RandomTree`` で木の形を指定するときに使用する列挙型クラスです。"""
160
161    random = enum.auto()
162    path = enum.auto()
163    star = enum.auto()
164
165
166class RandomTree:
167
168    @classmethod
169    def build(
170        cls,
171        n: int,
172        typ: RandomTreeType = RandomTreeType.random,
173        seed: Optional[int] = None,
174    ) -> list[tuple[int, int]]:
175        """ランダムな木を生成し、辺を返します。
176        :math:`O(n \\log{n})` です。
177
178        Args:
179            n (int): 頂点の数です。
180            typ (RandomTreeType, optional): 木の形です。 Defaults to RandomTreeType.random。
181            seed (Optional[int], optional): seed値です。 Defaults to None。
182
183        Returns:
184            list[tuple[int, int]]: 辺のリストです。辺のインデックスは 0-indexed です。
185        """
186        cls.rand = random.Random(seed)
187        edges = None
188        if typ == RandomTreeType.random:
189            edges = cls._build_random(n)
190        elif typ == RandomTreeType.path:
191            edges = cls._build_path(n)
192        elif typ == RandomTreeType.star:
193            edges = cls._build_star(n)
194        assert (
195            edges is not None
196        ), f"{cls.__class__.__name__}.build({typ}), typ is not defined."
197        cls.rand.shuffle(edges)
198        return edges
199
200    @classmethod
201    def _build_star(cls, n: int) -> list[tuple[int, int]]:
202        center = cls.rand.randrange(0, n)
203        edges = []
204        for i in range(n):
205            if i == center:
206                continue
207            if cls.rand.random() < 0.5:
208                edges.append((center, i))
209            else:
210                edges.append((i, center))
211        return edges
212
213    @classmethod
214    def _build_path(cls, n: int) -> list[tuple[int, int]]:
215        p = list(range(n))
216        cls.rand.shuffle(p)
217        edges = [
218            (p[i], p[i + 1]) if cls.rand.random() < 0.5 else (p[i + 1], p[i])
219            for i in range(n - 1)
220        ]
221        return edges
222
223    @classmethod
224    def _build_random(cls, n: int) -> list[tuple[int, int]]:
225        edges = []
226        D = [1] * n
227        A = [0] * (n - 2)
228        for i in range(n - 2):
229            v = cls.rand.randrange(0, n)
230            D[v] += 1
231            A[i] = v
232        avl: DeletableMinHeap[tuple[int, int]] = DeletableMinHeap(
233            (D[i], i) for i in range(n)
234        )
235        for a in A:
236            d, v = avl.pop_min()
237            assert d == 1
238            edges.append((v, a))
239            D[v] -= 1
240            avl.remove((D[a], a))
241            D[a] -= 1
242            if D[a] >= 1:
243                avl.push((D[a], a))
244        u = D.index(1)
245        D[u] -= 1
246        v = D.index(1)
247        D[v] -= 1
248        edges.append((u, v))
249        return edges

仕様

class RandomTree[source]

Bases: object

classmethod build(n: int, typ: RandomTreeType = RandomTreeType.random, seed: int | None = None) → list[tuple[int, int]][source]

ランダムな木を生成し、辺を返します。 \(O(n \log{n})\) です。

Parameters:

  • n (int) – 頂点の数です。

  • typ (RandomTreeType, optional) – 木の形です。 Defaults to RandomTreeType.random。

  • seed (Optional[int], optional) – seed値です。 Defaults to None。

Returns:

辺のリストです。辺のインデックスは 0-indexed です。

Return type:

list[tuple[int, int]]

class RandomTreeType(*values)[source]

Bases: Enum

RandomTree で木の形を指定するときに使用する列挙型クラスです。

path = 2
random = 1
star = 3