from copy import deepcopy
def test_all(f):
for i, data in enumerate(eval(f'test_data{f.__name__[0]}')):
exp = data["out"]
ans = f(*deepcopy(data["in"]))
result = "AC" if exp == ans else "WA"
print(f"{i+1} {result}: expected: {exp}, output: {ans}")
という訳で、上のコードを参考に隣(前 or 次)の順列を求める関数(ジェネレーター)を作ってみました。
from operator import lt, gt
def next_perm(a, rev=False):
a = a.copy()
op = gt if rev else lt
n = len(a)
while True:
for i in range(n - 2, -1, -1):
if op(a[i], a[i + 1]): break
else:
return None
for j in range(n - 1, i, -1):
if op(a[i], a[j]): break
a[i], a[j] = a[j], a[i]
a = a[:i + 1] + a[i + 1:][::-1]
yield a.copy()from itertools import permutations
# 昇順
a = list(range(5))
list(list(p) for p in permutations(a))[1:] == list(next_perm(a)) # 降順
a = list(range(5))
list(list(p) for p in permutations(a))[::-1][1:] == list(next_perm(a[::-1], rev=True))
a = [1, 2, 3, 3]
p1 = sorted([ list(p) for p in set(list(permutations(a)))])
p2 = [a] + list(next_perm(a))
p1 == p2for a, b in zip(p1, p2):
print(a, b)
class BIT:
def __init__(self, a):
if hasattr(a, "__iter__"):
self.size = len(a)
self.data = [0] * self.size
self.bit = [0] * (self.size + 1)
for i, j in enumerate(a):
self.add(i, j)
elif isinstance(a, int):
self.size = a
self.data = [0] * a
self.bit = [0] * (a + 1)
else:
raise TypeError()
def add(self, i, x):
self.data[i] += x
i += 1
while i <= self.size:
self.bit[i] += x
i += i & -i
def sum(self, i):
'''Get sum of data[:i]
'''
total = 0
while i > 0:
total += self.bit[i]
i -= i & -i
return total
def __getitem__(self, key):
return self.data[key]
def __setitem__(self, key, value):
self.add(key, value - self[key])
self.data[key] = value
def __len__(self):
return self.size
def __repr__(self):
return self.data.__repr__()
def __iter__(self):
return self.data.__iter__()
class Perm_n():
def __init__(self, items):
self.n = len(items)
# 階乗(0!, 1!, 2!, ..., n!)を保持
self.fa = [1]
# itemが辞書順で何番目(0-index)かを保持
self.order = {}
for i, item in enumerate(sorted(items), 1):
self.fa.append(self.fa[-1] * i)
self.order[item] = i - 1
def nth_perm(self, x):
S = list(self.order.keys())
ans = []
for i in range(self.n):
d, m = divmod(x, self.fa[self.n - 1 - i])
ans.append(S[d])
S = S[:d] + S[d+1:]
x = m
return ans
def id_of_perm(self, perm):
# i番目以降にi番目より小さい要素がいくつあるか
bit = BIT(self.n)
smaller = [None] * self.n
for i in range(self.n - 1, -1, -1):
o = self.order[perm[i]]
bit[o] = 1
smaller[i] = bit.sum(o)
ans = 0
for i in range(self.n):
ans += self.fa[self.n - 1 - i] * smaller[i]
return ansdef C(N, P):
out = StringIO()
perm = Perm_n(P)
print(*perm.nth_perm(perm.id_of_perm(P) - 1), end="", file=out)
return out.getvalue()test_all(C)
from math import gcd
def D(N, a):
gcd_a = a[0]
for ai in a[1:]:
gcd_a = gcd(gcd_a, ai)
cnt = 0
for ai in a:
ai //= gcd_a
while ai % 3 == 0:
ai //= 3
cnt += 1
while ai % 2 == 0:
ai //= 2
cnt += 1
if ai != 1: return -1
return cnttest_all(D)