| """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() |
| |