"""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()
Codebase Browser
llvm