from typing import List def cons(it, elem): yield from it yield elem class Cosets: def __init__(self, ngens, data=()): assert len(data) % ngens == 0, 'invalid length starting row' self.ngens = ngens self.data = list(data) self.len = len(data) // ngens def add_row(self): self.len += 1 self.data.extend(-1 for _ in range(self.ngens)) def put(self, idx, target): coset, gen = divmod(idx, self.ngens) self.data[idx] = target self.data[target * self.ngens + gen] = coset def get(self, idx): return self.data[idx] def __getitem__(self, key): coset, gen = key return self.data[coset * self.ngens + gen] def __setitem__(self, key, target): coset, gen = key self.data[coset * self.ngens + gen] = target self.data[target * self.ngens + gen] = coset def __len__(self): return self.len def __repr__(self): return '\n'.join( f'{i // self.ngens:>3} | ' + ' '.join( f'{e:>3}' for e in self.data[i:i + self.ngens]) for i in range(0, len(self.data), self.ngens)) class RelTable: def __init__(self, gens: List[int], mul: int): self.gens = gens self.mul = mul self.fam = [] self.gen = [] self.lst = [] def add_row(self): idx = len(self.fam) self.fam.append(-1) self.gen.append(-1) self.lst.append(-1) return idx class Group: def __init__(self, ngens, rels=()): self._mults = [[2] * ngens for _ in range(ngens)] for (f, t), m in rels: self._mults[f][t] = m self._mults[t][f] = m self.ngens = ngens @property def mults(self): for i in range(self.ngens - 1): for j in range(i + 1, self.ngens): yield ((i, j), self._mults[i][j]) def __setitem__(self, key, value): f, t = key self._mults[f][t] = value self._mults[t][f] = value def __mul__(self, other): assert isinstance(other, Group) off = self.ngens g = Group(self.ngens + other.ngens) for (i, j), m in self.mults: g[i, j] = m for (i, j), m in other.mults: g[off + i, off + j] = m return g def __pow__(self, p, modulo=None): if modulo is not None: raise NotImplemented assert isinstance(p, int), 'p must be an integer' assert p >= 0, 'p must be a nonnegative integer' g = Group(self.ngens * p) for (i, j), m in self.mults: for off in range(0, g.ngens, self.ngens): g[off + i, off + j] = m return g @classmethod def schlafli(cls, *mults): ngens = len(mults) + 1 return Group(ngens, (((i, i + 1), mult) for i, mult in enumerate(mults))) def solve(self, sub_gens=()): initial_row = [-1] * self.ngens for s in sub_gens: initial_row[s] = 0 cosets = Cosets(self.ngens, initial_row) rel_tables = [RelTable(*a) for a in self.mults] for rel in rel_tables: idx = rel.add_row() count = 0 for g in rel.gens: if cosets.get(g) == 0: count += 1 rel.fam[idx] = 0 rel.gen[idx] = 0 rel.lst[idx] = 0 if count == 1: rel.gen[idx] = -1 idx = 0 while True: while idx < len(cosets.data) and cosets.get(idx) >= 0: idx += 1 if idx == len(cosets.data): break coset, gen = divmod(idx, cosets.ngens) target = len(cosets) cosets.add_row() for rel in rel_tables: rel.add_row() facts = [(coset, gen)] while facts: coset, gen = facts.pop() cosets[coset, gen] = target for rel in rel_tables: if gen in rel.gens and rel.fam[target] == -1: rel.fam[target] = rel.fam[coset] rel.gen[target] = rel.gen[coset] + 1 if rel.gen[coset] < 0: rel.gen[target] -= 2 if rel.gen[target] == rel.mul: # forward learn lst = rel.lst[rel.fam[target]] gen_ = rel.gens[rel.gens[0] == gen] facts.append((lst, gen_)) elif rel.gen[target] == -rel.mul: # stationary learn gen_ = rel.gens[rel.gens[0] == gen] facts.append((target, gen_)) elif rel.gen[target] == rel.mul - 1: rel.lst[rel.fam[target]] = target facts.sort(reverse=True) for rel in rel_tables: if rel.fam[target] == -1: rel.fam[target] = target rel.gen[target] = 0 count = 0 for g in rel.gens: if cosets[target, g] == target: count += 1 if count == 1: rel.gen[target] = -1 return cosets