AtCoder Beginner Contest 222 E

https://atcoder.jp/contests/abc222/tasks/abc222_e

最短経路が辺 iを通る回数を d_iとすると、母関数 Gを考えて、
 \displaystyle G(x) = \prod_i{(x^{-d_i}+x^{d_i})}
この x^Kの係数を求めればよいです。 s = \sum_i{d_i}と置いて、
 \displaystyle G(x) = x^{-s}\prod_i{(1+x^{2d_i})}
 X = x^2として、 \displaystyle G_1(X) = \prod_i{(1+X^{d_i})}と置けば、
 G(x) = x^{-s}G_1(X)
なので、 G_1(X)を計算して X^{\frac{K+s}{2}}の係数を求めればよいです。

同じ dの重複度を dup(d)と書くと、
 \displaystyle G_1(X) = \prod_d{(1+X^d)^{dup(d)}}

これを最初分割統治法的に求めたら、すごく遅いんですね。ランダムな経路で計算したら14秒でした。各多項式が「疎な」多項式なので順に積を求める方が速いんですね。そのようにしたら1.1秒になりました。
さらに、あとに掛けるほうが項数が少ない方が速くなるはずなので、 dup(d)でソートして掛けると、0.8秒になりました。

以下は、PyPy2のコードです。

# coding: utf-8
# Red and Blue Tree

from collections import defaultdict, Counter
from itertools import count


#################### library ####################

def read_int():
    return int(raw_input())

def read_tuple():
    return tuple(map(int, raw_input().split()))

def read_list():
    return list(map(int, raw_input().split()))


#################### polynomial ####################

def poly_mul(f, g, d, D):
    L1, L2 = len(f), len(g)
    L3 = L1+(L2-1)*d
    h = [0] * L3
    for k in range(L2):
        for l in range(L1):
            h[l+k*d] += f[l]*g[k]
    return [ int(c % D) for c in h ]

def poly_pow(f, e, D):
    if e == 1:
        return f
    elif e%2 == 1:
        return poly_mul(f, poly_pow(f, e-1, D), D)
    else:
        g = poly_pow(f, e/2, D)
        return poly_mul(g, g, D)

def binomial(e):
    v = [1]
    c = 1
    for k in range(1, e+1):
        c = c * (e - k + 1) / k
        v.append(c)
    return v


#################### naive ####################


#################### process ####################

def read_input():
    N, M, K = read_tuple()
    As = read_list()
    edges = [ read_tuple() for _ in range(N-1) ]
    return (As, edges, K)

def make_tree(edges):
    graph = defaultdict(list)
    for U, V in edges:
        graph[U].append(V)
        graph[V].append(U)
    return graph

def shortest_path(v, w, graph):
    if v == w:
        return [v]
    
    a = [v]
    paths = { v: [v] }
    for d in count(1):
        b = []
        for v1 in a:
            for v2 in graph[v1]:
                if v2 in paths:
                    continue
                paths[v2] = paths[v1] + [v2]
                if v2 == w:
                    return paths[w]
                b.append(v2)
        a = b

def search_all_path(seq, graph):
    all_path = []
    for v, w in zip(seq, seq[1:]):
        path = shortest_path(v, w, graph)
        if all_path:
            all_path.extend(path[1:])
        else:
            all_path.extend(path)
    return all_path

def F(seq, graph, K):
    path = search_all_path(seq, graph)
    if len(path) % 2 == K % 2:
        return 0
    
    counter = Counter()
    for v1, v2 in zip(path, path[1:]):
        counter[(min(v1, v2), max(v1, v2))] += 1
    
    counter2 = Counter()
    for freq in counter.values():
        counter2[freq] += 1
    
    ens = counter2.items()
    s = sum(n for e, n in ens)
    f = [1]
    for d, e in sorted(ens, key=lambda (d, e): e, reverse=True):
        b = binomial(e)
        f = poly_mul(f, b, d, D)
    k = (len(path)-1+K)/2
    if 0 <= k < len(f):
        return f[k] * pow(2, len(graph) - s - 1, D) % D
    else:
        return 0


#################### main ####################

D = 998244353
seq, edges, K = read_input()
graph = make_tree(edges)
print F(seq, graph, K)

競プロ典型90問 002

https://atcoder.jp/contests/typical90/tasks/typical90_b

これは、wikipedia:カタラン数ですね。格子状の経路の数え方に対応していて、右に一つ移動すると'('、上に一つ移動すると')'です。
単純に1歩ずつ進むだけでOKです。

def F(N):
    def g(x, y):
        if (x, y) == (N/2, N/2):
            yield ''
        if x < N/2:
            for s in g(x+1, y):
                yield '(' + s
        if x > y:
            for s in g(x, y+1):
                yield ')' + s
    
    if N % 2 == 1:
        return
    
    for s in g(0, 0):
        print s

N = 20ならこれでも十分です。もう少し数を大きくしてみましょう。

$ echo 26 | time pypy typical002.py > /dev/null
1.32user 0.12system 0:01.46elapsed 99%CPU (0avgtext+0avgdata 56700maxresident)k
0inputs+0outputs (0major+21924minor)pagefaults 0swaps
$ echo 28 | time pypy typical002.py > /dev/null
5.07user 0.14system 0:05.24elapsed 99%CPU (0avgtext+0avgdata 57104maxresident)k
0inputs+0outputs (0major+34368minor)pagefaults 0swaps

どうやら26が限界のようです。
上の方法には非常に無駄があります。スタートを(0, 0)、ゴールを(N/2, N/2)として、例えば(3, 2)を通る経路は何通りかありますが、その度にそこから(N/2, N/2)までの経路を辿らなければいけません。かといって、メモ化だとやたらメモリを使うし。

分割統治法はできないでしょうか。すぐに思いつくのは、x + y = N/2で三角形を分割することです。

f:id:inamori:20211009142257p:plain

この斜めの線のどこかまで到達する経路を全部作って、それを反転してつなぎ合わせます。これならある程度無駄が少なくなって、メモリもそんなに使いません。
辞書順ということで、少しめんどうです。

def reverse(s):
    return ''.join(')' if c == '(' else '(' for c in s[::-1])

def F(N):
    def g(x, y, s):
        if x + y == N/2:
            yield (x, y), s
            return
        
        for v in g(x+1, y, s + '('):
            yield v
        if x > y:
            for v in g(x, y+1, s + ')'):
                yield v
    
    if N % 2 == 1:
        return
    
    dic = defaultdict(list)
    half_paths = list(g(0, 0, ''))
    for pt, s in half_paths:
        dic[pt].append(reverse(s))
    dic2 = dict((pt, sorted(ss)) for pt, ss in dic.items())
    
    counter = 0
    for pt, s1 in half_paths:
        for s2 in dic2[pt]:
            print s1 + s2
            counter += 1
echo 28 | time pypy typical002a.py > /dev/null
0.26user 0.12system 0:00.54elapsed 71%CPU (0avgtext+0avgdata 59664maxresident)k
0inputs+0outputs (0major+16019minor)pagefaults 0swaps
$ echo 30 | time pypy typical002a.py > /dev/null
1.03user 0.09system 0:01.13elapsed 98%CPU (0avgtext+0avgdata 59524maxresident)k
0inputs+0outputs (0major+20179minor)pagefaults 0swaps
$ echo 32 | time pypy typical002a.py > /dev/null
3.43user 0.42system 0:03.86elapsed 99%CPU (0avgtext+0avgdata 58940maxresident)k
0inputs+0outputs (0major+31366minor)pagefaults 0swaps

30まで行けるようになりました。
これ以上はうまくいきませんね。

競プロ典型90問 001

https://atcoder.jp/contests/typical90/tasks/typical90_a

話題になっていた問題です。
少し考えれば、スコア以上が成り立つかは簡単に求められるので、二分探索するだけですね。これで十分速いです。

以下は、PyPy2のコードです。

# coding: utf-8
# Yokan Party

from itertools import count


#################### library ####################

def read_int():
    return int(raw_input())

def read_tuple():
    return tuple(map(int, raw_input().split()))

def read_list():
    return list(map(int, raw_input().split()))


#################### naive ####################


#################### process ####################

def read_input():
    N, L = read_tuple()
    K = read_int()
    As = read_list()
    return (L, K, As)

def F(L, K, As):
    def is_cuttable(length):
        n = 0
        prev_pt = 0
        for A in As:
            if A - prev_pt >= length:
                n += 1
                if n == K+1:
                    return True
                prev_pt = A
        else:
            return L - prev_pt >= length and n == K
    
    def bin_search(first, last):
        if first == last - 1:
            return first
        
        mid = (first + last) / 2
        if is_cuttable(mid):
            return bin_search(mid, last)
        else:
            return bin_search(first, mid)
    
    return bin_search(1, L/(K+1)+1)


#################### main ####################

L, K, As = read_input()
print F(L, K, As)

AtCoder Beginner Contest 221 D

https://atcoder.jp/contests/abc221/tasks/abc221_d

分割統治法で、区間の集まりを[(first, last, num)]で表して、それをマージする、というコードを書いたら意外と長いコードになってしまって、max1462msec。
解説を見ると、ああって感じで、書いてみたら実質8行。でも、1216msecであまり変わりませんでした。最初に書いたコードは、考え方は汎用性があると思います。

以下は、PyPy2のコードです。

# coding: utf-8
# Online games

from collections import Counter


#################### library ####################

def read_int():
    return int(raw_input())

def read_tuple():
    return tuple(map(int, raw_input().split()))

def read_list():
    return list(map(int, raw_input().split()))

def digits(n):
    while n:
        n, d = divmod(n, 10)
        yield d


#################### naive ####################

def F_naive(intervals):
    c = Counter()
    first = min(A for A, B in intervals)
    last = max(A+B for A, B in intervals)
    for d in range(first, last):
        n = sum(1 for A, B in intervals if A <= d < A+B)
        c[n] += 1
    return [ c[n] for n in range(1, len(intervals) + 1) ]


#################### process ####################

def read_input():
    N = read_int()
    intervals = [ read_tuple() for _ in range(N) ]
    return intervals

import random
def random_intervals(first, last):
    A = random.randrange(first, last)
    B = random.randrange(1, last - first)
    return (A, B)

def F(intervals):
    def g(ints):
        if len(ints) == 1:
            A, B = ints[0]
            yield (A, A+B, 1)
            return
        
        mid = len(ints) / 2
        ints1 = list(g(ints[:mid]))
        ints2 = list(g(ints[mid:]))
        k, l = 0, 0
        while k < len(ints1) and l < len(ints2):
            A1, B1, n1 = ints1[k]
            A2, B2, n2 = ints2[l]
            if B1 == B2:
                if A1 == A2:
                    yield (A1, B1, n1+n2)
                elif A1 < A2:
                    yield (A1, A2, n1)
                    yield (A2, B1, n1+n2)
                else:
                    yield (A2, A1, n2)
                    yield (A1, B1, n1+n2)
                k += 1
                l += 1
            elif B1 < B2:
                if B1 <= A2:
                    yield ints1[k]
                elif A1 == A2:
                    yield (A1, B1, n1+n2)
                    ints2[l] = (B1, B2, n2)
                elif A1 < A2:
                    yield (A1, A2, n1)
                    yield (A2, B1, n1+n2)
                    ints2[l] = (B1, B2, n2)
                else:
                    yield (A2, A1, n2)
                    yield (A1, B1, n1+n2)
                    ints2[l] = (B1, B2, n2)
                k += 1
            else:
                if B2 <= A1:
                    yield ints2[l]
                elif A1 == A2:
                    yield (A1, B2, n1+n2)
                    ints1[k] = (B2, B1, n1)
                elif A1 < A2:
                    yield (A1, A2, n1)
                    yield (A2, B2, n1+n2)
                    ints1[k] = (B2, B1, n1)
                else:
                    yield (A2, A1, n2)
                    yield (A1, B2, n1+n2)
                    ints1[k] = (B2, B1, n1)
                l += 1
        
        for int1 in ints1[k:]:
            yield int1
        
        for int2 in ints2[l:]:
            yield int2
    
    c = Counter()
    for first, last, n in g(intervals):
        c[n] += last - first
    return [ c[n] for n in range(1, len(intervals) + 1) ]


#################### main ####################

intervals = read_input()
#intervals = [ random_intervals(1, 100) for _ in range(3) ]
#print F_naive(intervals)
v = F(intervals)
print ' '.join(map(str, v))

AtCoder Beginner Contest 215 E

https://atcoder.jp/contests/abc215/tasks/abc215_e

ただのDPで、PyPy2だとあっさり通りました。
で、Python3で試したら、24ケース中13ケースでTLE。案の定Pythonでは遅いです。
手元でも試したら、実行時間は元の文字列の長さにだいたい比例するので、長さ少し変えて、

length Python2 PyPy2
1000 1.92s 0.25s
2000 3.83s 0.44s
4000 7.80s 0.84s

だいたいPyPy2のほうが9倍くらい速いですね。長いリストを使うとだいたいこんな感じです。
たいていの場合PyPyのほうが速いので、素直にPyPyを使いましょう。

AtCoder Regular Contest 118 D

https://atcoder.jp/contests/arc118/tasks/arc118_d

面白かったので解説してみます。

P = 13, a = 4, b = 5で考えてみます。
まず、グラフを描いてみましょう。

a倍していくと、
1 -> 4 -> 3 -> 12 -> 9 -> 10 -> 1
と周期6になります。
もう1つサークルがあります。
2 -> 8 -> 6 -> 11 -> 5 -> 7 -> 2

b倍していくと周期4が3つあります。
1 -> 5 -> 12 -> 8 -> 1
2 -> 10 -> 11 -> 3 -> 2
4 -> 7 -> 9 -> 6 -> 4

この2つの周期のサークルが絡み合っているので難しそうですよね。

さて、素数pに対して1 ~ p-1はaを適当に選ぶと周期p-1のサークルになります。p = 13に対してa = 6とすると、
1 -> 6 -> 10 -> 8 -> 9 -> 2 -> 12 -> 7 -> 3 -> 5 -> 4 -> 11 -> 1
周期12になります。
 g = 6と置くと、 4 \equiv g^{10}なので、a倍していくサークルは、
 1 \equiv g^0 \rightarrow g^{10} \rightarrow g^{20} \rightarrow g^{30} \rightarrow g^{40} \rightarrow g^{50} \rightarrow g^{60} \equiv 1
 2 \equiv g^5 \rightarrow g^{15} \rightarrow g^{25} \rightarrow g^{35} \rightarrow g^{45} \rightarrow g^{55} \rightarrow g^{65} \equiv 2

 5 \equiv g^9なので、b倍していくサークルは、
 1 \equiv g^0 \rightarrow g^9 \rightarrow g^{18} \rightarrow g^{27} \rightarrow g^{36} \equiv 1
 2 \equiv g^5 \rightarrow g^{14} \rightarrow g^{23} \rightarrow g^{32} \rightarrow g^{41} \equiv 2
 4 \equiv g^{10} \rightarrow g^{19} \rightarrow g^{28} \rightarrow g^{37} \rightarrow g^{46} \equiv 4

指数だけ見ましょう。そうすると、掛け算は足し算になって、12を法とする剰余を考えます。そして、指数を並べると、a倍を横、b倍を縦に並べて、

36 46 56 66 76 86 96
27 37 47 57 67 77 87
18 28 38 48 58 68 78
 9 19 29 39 49 59 69
 0 10 20 30 40 50 60

一番左と右、一番下と上は同じ数を表しています。つまり、上下左右に無限に伸びています。トーラスですね。
グラフの辺は、それぞれ数とその上下左右の数を結びます。
パスは、下の2行だけ考えて、

 9 → 19    29 → 39    49 → 59
 ↑    ↓    ↑     ↓    ↑     ↓
 0    10 → 20    30 → 40    50 → 60

とすればよいです。この方法は、横の周期が偶数だと成り立ちます。
ただし、b = 10だとパスはありません。なぜなら、 10 \equiv g^2なので、1にどうaとbを掛けても指数は偶数になってしまい、奇数は表せないからです。
 a \equiv g^{e_a}とすると、周期は p_a = (P-1)/\gcd{(P-1, e_a)}となります。逆に周期が分かっていたら、 \gcd{(P-1, e_a)} = (P-1)/p_aとなります。bの方の周期を p_bとすると、

 \gcd{(P-1, \gcd{(e_a, e_b)})} = \gcd{((P-1)/p_a, (P-1)/p_b)}

となります。これが1でないならパスは無しです。
周期はaを次々と掛けていけば得られると思いますが、周期はP-1の約数なので、約数の小さい順にべき乗の計算をすればよいでしょう。

最後に、横の周期を必ず偶数にできるかです。 p_a, p_bのどちらかが偶数になればいいので、どちらも奇数があり得るか考えます。
P-1はP=2以外では偶数です(P=2のときは、1 → 1というパスでOKです)。 p_a, p_bが共に奇数なら、 \gcd{((P-1)/p_a, (P-1)/p_b)}は偶数になります。

以下は、PyPy2のコードです。

# coding: utf-8
# Hamiltonian Cycle

from itertools import *
from fractions import gcd
import sys


#################### library ####################

def factorize_naive(n):
	def f(n, p0):
		if n <= 1:
			return ()
		
		for p in takewhile(lambda p: p * p <= n, count(p0)):
			if n % p == 0:
				return (p,) + f(n / p, p)
		else:
			return (n,)
	
	return list((k, len(list(v))) for k, v in groupby(f(n, 2)))

def divisors_f(fs):
	if len(fs) == 1:
		p, e = fs[0]
		for e1 in range(e+1):
			yield p**e1
	else:
		mid = len(fs) / 2
		ds1 = list(divisors_f(fs[:mid]))
		ds2 = list(divisors_f(fs[mid:]))
		for d1 in ds1:
			for d2 in ds2:
				yield d1 * d2

def period(e, p):
	fs = factorize_naive(p - 1)
	ds = sorted(divisors_f(fs))
	for d in ds:
		if pow(e, d, p) == 1:
			return d


#################### process ####################

def read_input():
	P, a, b = tuple(map(int, raw_input().split()))
	return (P, a, b)

def is_valid(path, P, a, b):
	if len(path) != P:
		return False
	
	s = set(path[:-1])
	if len(s) != P - 1:
		return False
	
	for v1, v2 in zip(path, path[1:]):
		if not (a * v1 % P == v2 or a * v2 % P == v1 or
				b * v1 % P == v2 or b * v2 % P == v1):
			return False
	else:
		return True

def Hamiltonian_Cycle(P, a, b):
	def generate_path(p1, e1, p2, e2):
		for i in range(p1):
			if i % 2 == 0:
				for j in range(p2):
					yield pow(e1, i, P) * pow(e2, j, P) % P
			else:
				for j in range(p2 - 1, -1, -1):
					yield pow(e1, i, P) * pow(e2, j, P) % P
		yield 1
	
	if P == 2:
		return [1, 1]
	
	pa = period(a, P)
	pb = period(b, P)
	if gcd((P - 1) / pa, (P - 1) / pb) != 1:
		return []
	
	if pa % 2 == 0:
		return list(generate_path(pa, a, (P - 1) / pa, b))
	else:
		return list(generate_path(pb, b, (P - 1) / pb, a))


#################### main ####################

P, a, b = read_input()
cycle = Hamiltonian_Cycle(P, a, b)
#print(is_valid(cycle, P, a, b))
if cycle:
	print('Yes')
	print(' '.join(map(str, cycle)))
else:
	print('No')