| """Utilities for enumeration of finite and countably infinite sets. |
| """ |
| from __future__ import absolute_import, division, print_function |
| |
| ### |
| # Countable iteration |
| |
| # Simplifies some calculations |
| class Aleph0(int): |
| _singleton = None |
| |
| def __new__(type): |
| if type._singleton is None: |
| type._singleton = int.__new__(type) |
| return type._singleton |
| |
| def __repr__(self): |
| return "<aleph0>" |
| |
| def __str__(self): |
| return "inf" |
| |
| def __cmp__(self, b): |
| return 1 |
| |
| def __sub__(self, b): |
| raise ValueError("Cannot subtract aleph0") |
| |
| __rsub__ = __sub__ |
| |
| def __add__(self, b): |
| return self |
| |
| __radd__ = __add__ |
| |
| def __mul__(self, b): |
| if b == 0: |
| return b |
| return self |
| |
| __rmul__ = __mul__ |
| |
| def __floordiv__(self, b): |
| if b == 0: |
| raise ZeroDivisionError |
| return self |
| |
| __rfloordiv__ = __floordiv__ |
| __truediv__ = __floordiv__ |
| __rtuediv__ = __floordiv__ |
| __div__ = __floordiv__ |
| __rdiv__ = __floordiv__ |
| |
| def __pow__(self, b): |
| if b == 0: |
| return 1 |
| return self |
| |
| |
| aleph0 = Aleph0() |
| |
| |
| def base(line): |
| return line * (line + 1) // 2 |
| |
| |
| def pairToN(pair): |
| x, y = pair |
| line, index = x + y, y |
| return base(line) + index |
| |
| |
| def getNthPairInfo(N): |
| # Avoid various singularities |
| if N == 0: |
| return (0, 0) |
| |
| # Gallop to find bounds for line |
| line = 1 |
| next = 2 |
| while base(next) <= N: |
| line = next |
| next = line << 1 |
| |
| # Binary search for starting line |
| lo = line |
| hi = line << 1 |
| while lo + 1 != hi: |
| # assert base(lo) <= N < base(hi) |
| mid = (lo + hi) >> 1 |
| if base(mid) <= N: |
| lo = mid |
| else: |
| hi = mid |
| |
| line = lo |
| return line, N - base(line) |
| |
| |
| def getNthPair(N): |
| line, index = getNthPairInfo(N) |
| return (line - index, index) |
| |
| |
| def getNthPairBounded(N, W=aleph0, H=aleph0, useDivmod=False): |
| """getNthPairBounded(N, W, H) -> (x, y) |
| |
| Return the N-th pair such that 0 <= x < W and 0 <= y < H.""" |
| |
| if W <= 0 or H <= 0: |
| raise ValueError("Invalid bounds") |
| elif N >= W * H: |
| raise ValueError("Invalid input (out of bounds)") |
| |
| # Simple case... |
| if W is aleph0 and H is aleph0: |
| return getNthPair(N) |
| |
| # Otherwise simplify by assuming W < H |
| if H < W: |
| x, y = getNthPairBounded(N, H, W, useDivmod=useDivmod) |
| return y, x |
| |
| if useDivmod: |
| return N % W, N // W |
| else: |
| # Conceptually we want to slide a diagonal line across a |
| # rectangle. This gives more interesting results for large |
| # bounds than using divmod. |
| |
| # If in lower left, just return as usual |
| cornerSize = base(W) |
| if N < cornerSize: |
| return getNthPair(N) |
| |
| # Otherwise if in upper right, subtract from corner |
| if H is not aleph0: |
| M = W * H - N - 1 |
| if M < cornerSize: |
| x, y = getNthPair(M) |
| return (W - 1 - x, H - 1 - y) |
| |
| # Otherwise, compile line and index from number of times we |
| # wrap. |
| N = N - cornerSize |
| index, offset = N % W, N // W |
| # p = (W-1, 1+offset) + (-1,1)*index |
| return (W - 1 - index, 1 + offset + index) |
| |
| |
| def getNthPairBoundedChecked( |
| N, W=aleph0, H=aleph0, useDivmod=False, GNP=getNthPairBounded |
| ): |
| x, y = GNP(N, W, H, useDivmod) |
| assert 0 <= x < W and 0 <= y < H |
| return x, y |
| |
| |
| def getNthNTuple(N, W, H=aleph0, useLeftToRight=False): |
| """getNthNTuple(N, W, H) -> (x_0, x_1, ..., x_W) |
| |
| Return the N-th W-tuple, where for 0 <= x_i < H.""" |
| |
| if useLeftToRight: |
| elts = [None] * W |
| for i in range(W): |
| elts[i], N = getNthPairBounded(N, H) |
| return tuple(elts) |
| else: |
| if W == 0: |
| return () |
| elif W == 1: |
| return (N,) |
| elif W == 2: |
| return getNthPairBounded(N, H, H) |
| else: |
| LW, RW = W // 2, W - (W // 2) |
| L, R = getNthPairBounded(N, H**LW, H**RW) |
| return getNthNTuple( |
| L, LW, H=H, useLeftToRight=useLeftToRight |
| ) + getNthNTuple(R, RW, H=H, useLeftToRight=useLeftToRight) |
| |
| |
| def getNthNTupleChecked(N, W, H=aleph0, useLeftToRight=False, GNT=getNthNTuple): |
| t = GNT(N, W, H, useLeftToRight) |
| assert len(t) == W |
| for i in t: |
| assert i < H |
| return t |
| |
| |
| def getNthTuple( |
| N, maxSize=aleph0, maxElement=aleph0, useDivmod=False, useLeftToRight=False |
| ): |
| """getNthTuple(N, maxSize, maxElement) -> x |
| |
| Return the N-th tuple where len(x) < maxSize and for y in x, 0 <= |
| y < maxElement.""" |
| |
| # All zero sized tuples are isomorphic, don't ya know. |
| if N == 0: |
| return () |
| N -= 1 |
| if maxElement is not aleph0: |
| if maxSize is aleph0: |
| raise NotImplementedError("Max element size without max size unhandled") |
| bounds = [maxElement**i for i in range(1, maxSize + 1)] |
| S, M = getNthPairVariableBounds(N, bounds) |
| else: |
| S, M = getNthPairBounded(N, maxSize, useDivmod=useDivmod) |
| return getNthNTuple(M, S + 1, maxElement, useLeftToRight=useLeftToRight) |
| |
| |
| def getNthTupleChecked( |
| N, |
| maxSize=aleph0, |
| maxElement=aleph0, |
| useDivmod=False, |
| useLeftToRight=False, |
| GNT=getNthTuple, |
| ): |
| # FIXME: maxsize is inclusive |
| t = GNT(N, maxSize, maxElement, useDivmod, useLeftToRight) |
| assert len(t) <= maxSize |
| for i in t: |
| assert i < maxElement |
| return t |
| |
| |
| def getNthPairVariableBounds(N, bounds): |
| """getNthPairVariableBounds(N, bounds) -> (x, y) |
| |
| Given a finite list of bounds (which may be finite or aleph0), |
| return the N-th pair such that 0 <= x < len(bounds) and 0 <= y < |
| bounds[x].""" |
| |
| if not bounds: |
| raise ValueError("Invalid bounds") |
| if not (0 <= N < sum(bounds)): |
| raise ValueError("Invalid input (out of bounds)") |
| |
| level = 0 |
| active = list(range(len(bounds))) |
| active.sort(key=lambda i: bounds[i]) |
| prevLevel = 0 |
| for i, index in enumerate(active): |
| level = bounds[index] |
| W = len(active) - i |
| if level is aleph0: |
| H = aleph0 |
| else: |
| H = level - prevLevel |
| levelSize = W * H |
| if N < levelSize: # Found the level |
| idelta, delta = getNthPairBounded(N, W, H) |
| return active[i + idelta], prevLevel + delta |
| else: |
| N -= levelSize |
| prevLevel = level |
| else: |
| raise RuntimError("Unexpected loop completion") |
| |
| |
| def getNthPairVariableBoundsChecked(N, bounds, GNVP=getNthPairVariableBounds): |
| x, y = GNVP(N, bounds) |
| assert 0 <= x < len(bounds) and 0 <= y < bounds[x] |
| return (x, y) |
| |
| |
| ### |
| |
| |
| def testPairs(): |
| W = 3 |
| H = 6 |
| a = [[" " for x in range(10)] for y in range(10)] |
| b = [[" " for x in range(10)] for y in range(10)] |
| for i in range(min(W * H, 40)): |
| x, y = getNthPairBounded(i, W, H) |
| x2, y2 = getNthPairBounded(i, W, H, useDivmod=True) |
| print(i, (x, y), (x2, y2)) |
| a[y][x] = "%2d" % i |
| b[y2][x2] = "%2d" % i |
| |
| print("-- a --") |
| for ln in a[::-1]: |
| if "".join(ln).strip(): |
| print(" ".join(ln)) |
| print("-- b --") |
| for ln in b[::-1]: |
| if "".join(ln).strip(): |
| print(" ".join(ln)) |
| |
| |
| def testPairsVB(): |
| bounds = [2, 2, 4, aleph0, 5, aleph0] |
| a = [[" " for x in range(15)] for y in range(15)] |
| b = [[" " for x in range(15)] for y in range(15)] |
| for i in range(min(sum(bounds), 40)): |
| x, y = getNthPairVariableBounds(i, bounds) |
| print(i, (x, y)) |
| a[y][x] = "%2d" % i |
| |
| print("-- a --") |
| for ln in a[::-1]: |
| if "".join(ln).strip(): |
| print(" ".join(ln)) |
| |
| |
| ### |
| |
| # Toggle to use checked versions of enumeration routines. |
| if False: |
| getNthPairVariableBounds = getNthPairVariableBoundsChecked |
| getNthPairBounded = getNthPairBoundedChecked |
| getNthNTuple = getNthNTupleChecked |
| getNthTuple = getNthTupleChecked |
| |
| if __name__ == "__main__": |
| testPairs() |
| |
| testPairsVB() |