Pythonによる螺旋本・AOJの解答（ALDS1 #1~#6）

この記事では、螺旋本・AOJのALDS1コース（アルゴリズムとデータ構造入門）のトピック#1〜#6にPythonで解答を与えます。

• 追加情報：この記事の続編Pythonによる螺旋本・AOJの解答（ALDS1 #7~#12）も書きました

AOJ とは

AOJ (Aizu Online Judge) は会津大学の渡部有隆准教授が開発したプログラミング自習サービスで、用意された問題に対してコードを入力すると自動採点してくれます。近年AOJ 2.0がリリースされ、UIなどがリニューアルされています。

このAOJの問題を解説していく本

• 渡部 有隆, Ozy（協力), 秋葉 拓哉（協力）『プログラミングコンテスト攻略のためのアルゴリズムとデータ構造』（表紙に螺旋の絵が描いてあるため、螺旋本とよく呼ばれています。AOJ本とも呼ばれています。名前が長すぎることからTLE本と呼ばれることもあるようです）

は、AtCoderといった競技プログラミング対策の自習のための本として、

• 秋葉 拓哉, 岩田 陽一, 北川 宜稔『プログラミングコンテストチャレンジブック [第2版]』（こちらは表紙に蟻の絵が描いてあるため、蟻本とよく呼ばれています）

に並んで人気がある本です。

コースは複数あるのですが、今回はALDS1コース（アルゴリズムとデータ構造入門）の問題を解きます。ALDS1には#1から#15までのトピックがあり、各トピックにはAからDまでの4問があります。今回は、トピック#1〜#6までの24問をPythonで解きます。以降のトピックに関しても、そのうちPythonによる解答を与えたいと思っています。

参考文献

• 学習する天然ニューラルネット　螺旋本をPythonで解く

AOJのPythonによる解答は既にいくつかありますが、とくにこの記事は一部参考にさせて頂きました。ありがとうございます。

目次

• トピック #1 入門
• トピック #2 初等的ソート
• トピック #3 基本データ構造
• トピック #4 探索
• トピック #5 分割統治法
• トピック #6 ソート

解答

トピック #1 入門

ALDS1_1_A: Insertion Sort

N = int(input())
A = [*map(int, input().split())]

for i in range(N):
    v = A[i]
    j = i-1
    while j >= 0 and A[j] > v:
        A[j+1] = A[j]
        j -= 1
    A[j+1] = v
    print(*A)

• 挿入ソートは、既にソート済みの部分配列へ要素を挿入していく安定なソート。最悪計算量は$\Theta(n^2)$と悪い。平均計算量も$\Theta(n^2)$のまま

ALDS1_1_B: Greatest Common Divisor

問題のコンセプトを無視してmathモジュールを使うと、

import math
print(math.gcd(*map(int, input().split())))

コンセプト通りに書けば、

x, y = map(int, input().split())
while y:
  x, y = y, x%y
print(x)

• いわゆるEuclidの互除法。$GCD(x,y)=GCD(y,x \text{ mod } y)$を用いて、$x,y$を幾何級数的に小さくしていく
• 大小関係が逆 $x < y$ なら、$x$と$y$は勝手に交換される

ALDS1_1_C: Prime Numbers

いわゆるEratosthenesの篩（ふるい）を使った自然なコードは

N = 10**8
primes = range(2, N+1)
for d in range(2, int(N**.5)+1):
  primes = [x for x in primes if x % d or x == d]

count = 0
for _ in range(int(input())):
    if int(input()) in primes:
        count += 1
print(count)

だろうけど、MLEやTLEになってしまう。入力される$2\le n \le 10^8$の整数$n$の全てに対応できるように$10^8$までの素数表を作っている部分が問題と思われる。

• 入力される整数の最大値を$N=10^8$とすると、素数の逆数の和が$O(\log^2 N)$であることなどから素数表を作る部分の計算量は$O(N \log^2 N)$となるそうで、計算量上限の目安の$O(10^8)$にひっかかっている

コードが少し複雑になるものの、素数表にある数で割れるかどうかを調べることによって、$10^4=\sqrt{10^8}$までの素数表でも$10^8$までの整数に対応でき、ACになる：

N = 10**4
primes = range(2, N+1)
for d in range(2, int(N**.5)+1):
  primes = [x for x in primes if x % d or x == d]

count = 0
for _ in range(int(input())):
    x = int(input())
    if x <= N:
        count += (x in primes)
    else:
        for prime in primes:
            if x % prime == 0:
                break
        else:
            count += 1
print(count)

• これで素数表を作る部分の計算量は$O(\sqrt{N} \log^2 N)$に減るはずで、計算量上限の目安$O(10^8)$を超えなさそう

あるいは、Eratosthenesの篩によって素数表を作る方針を捨てて、整数が入力される度に素数判定をしてもACになる。コードもシンプルになる：

def is_prime(x):
    for d in range(2, int(x**.5)+1):
        if x%d == 0:
            return False
    return True

n = int(input())
print(sum(is_prime(int(input())) for _ in range(n)))

• この場合、計算量は$O(n \sqrt{N})$となるはずで危ういが、実際は先ほどよりわずかに遅くなるくらいで通る

• 1から始まる自然数は、1か素数か合成数かの３つに分類される。関数is_prime(x)は実際は合成数でないことを判定しているだけなので、もし1が来るとTrueを返してしまうが、入力される自然数が2以上の問題なのでこれで十分

ALDS1_1_D: Maximum Profit

R_min, dR_max = float("inf"), -float("inf")
for _ in range(int(input())):
    R = int(input())
    dR_max = max(dR_max, R-R_min)
    R_min = min(R_min, R)
print(dR_max)

• 最小値R_minと差の最大値dR_maxの２変数さえ保持しておけば、入力を逐次的に処理できる
• 空集合におけるminとmaxをそれぞれ$+∞$, $-∞$と解釈することで、場合分けが減って綺麗になる

トピック #2 初等的ソート

ALDS1_2_A: Bubble Sort

N = int(input())
A = [*map(int, input().split())]

inv = 0
for i in range(N):
  for j in reversed(range(i+1, N)):
    if A[j] < A[j-1]:
      A[j-1], A[j] = A[j], A[j-1]
      inv += 1
print(*A)
print(inv)

• バブルソートは、右端から隣同士を比べて小さい方を左に動かし続けることで、既にソート済みの左側の部分配列を右に拡大させていく安定なソート。最悪計算量は$\Theta (n^2)$と悪い。平均計算量も$\Theta (n^2)$のままだそうだ

ALDS1_2_B: Selection Sort

N = int(input())
A = [*map(int, input().split())]

count = 0
for i in range(N-1):
  j_min = A.index(min(A[i:]), i)
  if A[i] > A[j_min]:
    A[i], A[j_min] = A[j_min], A[i]
    count += 1
print(*A)
print(count)

• 選択ソートは、ソート済みではない部分配列から最小値を抜き出し続けるソート。安定では無く、最悪計算量は$\Theta(n^2)$と悪い。最良計算量も$\Theta(n^2)$なので、平均計算量も$\Theta(n^2)$だ

ALDS1_2_C: Stable Sort

N = int(input())
A = input().split()

def bubble_sort(A, N):
  for i in range(N):
    for j in reversed(range(i+1, N)):
      if A[j][1] < A[j-1][1]:
        A[j-1], A[j] = A[j], A[j-1]

def selection_sort(A, N):
  for i in range(N-1):
    j_min = A.index(min(A[i:], key=lambda a:a[1]), i)
    A[i], A[j_min] = A[j_min], A[i]

A_bub = A[:]
bubble_sort(A_bub, N)
print(*A_bub)
print("Stable")

selection_sort(A, N)
print(*A)
print("Stable" if A == A_bub else "Not stable")

• バブルソートは安定だが、選択ソートは不安定
• 数字が1桁なので、intにしなくてもstrのまま比較できる

ALDS1_2_D: Shell Sort

N = int(input())
A = [int(input()) for _ in range(N)]

def insertion_sort(A, N, step):
  count = 0
  for i in range(step,N):
      v = A[i]
      j = i - step
      while j >= 0 and A[j]>v:
          A[j+step] = A[j]
          j -= step
          count += 1
      A[j+step] = v
  return count

m = int.bit_length(N)
print(m)
G = [N//(2**i) for i in range(m)]
print(*G)
print(sum(insertion_sort(A, N, G[i]) for i in range(m)))
print(*A, sep="\n")

• Shellソートは一定の距離離れた要素同士だけで挿入ソートし、その距離を徐々に狭めていく不安定なソート。最終的に距離は1になり、挿入ソートに帰着する

• ほぼ揃った入力に対しては挿入ソートは高速なので、Shellソートはただの挿入ソートよりも高速になる。距離をどのような単調列にするか次第で計算量のオーダーが変わり、現在知られている最良の最悪計算量は$\Theta(N \log^2 N)$のようで、平均計算量についてはさらによくわかっていないようだ

• この解答で使った単調列$\lfloor N/2^i \rfloor$は1959年にShellがShellソートを提案したときに用いていたもので、最悪計算量は$\Theta(n^2)$と悪い

トピック #3 基本データ構造

ALDS1_3_A: Stack

Pythonでは、スタックの実装には単にリストを使えば良い：

ops = {"+":lambda x,y: y+x, "-":lambda x,y: y-x, "*":lambda x,y: y*x}
stack = []
for a in input().split():
    stack.append(ops[a](stack.pop(), stack.pop()) if a in ops else int(a))
print(stack.pop())

eval()を使って文字列をPythonの式として実行して計算する方法もある：

stack = []
for a in input().split():
  if a.isdigit():
    stack.append(a)
  else:
    stack.append(str(eval("".join(reversed([stack.pop(), a, stack.pop()])))))
print(stack.pop())

ALDS1_3_B: Queue

from collections import deque

n, q = map(int, input().split())
procs = deque([name, int(time)] for name, time in [input().split() for _ in range(n)])

tot_time = 0
while procs:
  name, time = procs.popleft()
  if time <= q:
    tot_time += time
    print(name, tot_time)
  else:
    tot_time += q
    procs.append([name, time - q])

• Pythonのリストは末尾への追加や削除なら$O(1)$でできる一方で、先頭の値の追加や削除には$O(n)$かかってしまうため、両端での操作が必要なキューや両側キューとしては使いづらい。collectionsモジュールのdequeなら全て$O(1)$でできるので、リストよりdequeを使うのが良い

• 一方、両端では無い内部の要素へのランダムアクセスではdequeよりリストの方が速いようだ

ALDS1_3_C: Doubly Linked List

自然なコードは次のようなものだろうけど、TLEになってしまう：

from collections import deque

ops = {"insert": deque.appendleft,
       "delete": lambda deq, x: deq.remove(x) if x in deq else None,
       "deleteFirst": deque.popleft,
       "deleteLast": deque.pop}

deq = deque()
for _ in range(int(input())):
  op_name, *key = input().split()
  ops[op_name](deq, *key)
print(*deq)

どうやらinput().split()の計算が重いらしく、split()せずに処理できるように書き直すと可読性は下がるもののなんとか通る：

from collections import deque

ops = {"insert ": deque.appendleft,
       "delete ": lambda deq, x: deq.remove(x) if x in deq else None,
       "deleteF": deque.popleft,
       "deleteL": deque.pop}

deq = deque()
for _ in range(int(input())):
  op = input()
  if op[6] == " ":
    ops[op[:7]](deq, op[7:])
  else:
    ops[op[:7]](deq)
print(*deq)

ALDS1_3_D: Areas on the Cross-Section Diagram

downs, ponds = [], []
for i, s in enumerate(input()):
  if s == "\\":
    downs.append(i)
  elif s == "/" and downs:
    i_down = downs.pop()
    area = i - i_down
    while ponds and ponds[-1][0] > i_down:
      area += ponds.pop()[1]
    ponds.append([i_down, area])
print(sum(p[1] for p in ponds))
print(len(ponds), *(p[1] for p in ponds))

• downsはまだ水たまりになっていない各下り坂の位置i_downのスタックで、pondsは各水たまりの開始位置i_downと面積areaのスタック

トピック #4 探索

ALDS1_4_A: Linear Search

_, S = input(), input().split()
_, T = input(), input().split()
print(sum(t in S for t in T))

ALDS1_4_B: Binary Search

問題のコンセプトを無視してbisectモジュールを使うと、

import bisect

_, S = input(), [*map(int, input().split())]
_, T = input(), [*map(int, input().split())]

def binary_search(S, t):
    return t == S[bisect.bisect_left(S, t)]

print(sum(binary_search(S, t) for t in T))

使わずに自分で書くと、

_, S = input(), [*map(int, input().split())]
_, T = input(), [*map(int, input().split())]

def binary_search(S, t):
  left, right = 0, len(S)
  while left < right:
    mid = (left + right) // 2
    if S[mid] == t:
      return True
    if S[mid] > t:
      right = mid
    else:
      left = mid + 1
  return False

print(sum(binary_search(S, t) for t in T))

ALDS1_4_C: Dictionary

d = set()
for _ in range(int(input())):
  command, s = input().split()
  if command == "find":
    print("yes" if s in d else "no")
  else:
    d.add(s)

ALDS1_4_D: Allocation

n, k = map(int, input().split())
w = [int(input()) for _ in range(n)]

def n_loads(P):
  num = 0
  for _ in range(k):
    w_sum = 0
    while num < n and w_sum + w[num] <= P:
      w_sum += w[num]
      num += 1
  return num

P_lb, P_min = 0, max(w)*(n//k + 1)
while P_min - P_lb > 1:
  P_mid = (P_lb + P_min) // 2
  if n_loads(P_mid) < n:
    P_lb = P_mid
  else:
    P_min = P_mid
print(P_min)

• このように単調列の値が配列として予め得られておらず、毎回値を計算する必要があるような、言わばブラックボックス単調関数値の探索の場合の二分探索ではbisectモジュールが使えない

トピック #5 分割統治法

ALDS1_5_A: Exhaustive Search

分割統治法で解くという問題のコンセプトを無視してbit全探索でも解ける：

import itertools as it

n, A = int(input()), [*map(int, input().split())]
yes_nums = set()
for bools in it.product([True, False], repeat=n):
  yes_nums.add(sum(A[i] for i in range(n) if bools[i]))

_ = input()
for m in map(int, input().split()):
  print("yes" if m in yes_nums else "no")

• q個の各入力に対してAの要素の和で作れるかどうかを判定すると、$O(q2^n)$の計算量がかかってしまう。Aの要素の和として作れる数の表を予め作っておけば、計算量は$O(2^n+q)$で済む

既に作れた数を憶えておきながら解けば、より速い：

_ = input()
yes_nums = {0}
for x in map(int, input().split()):
  for y in [*yes_nums]:
    yes_nums.add(x + y)

_ = input()
for m in map(int, input().split()):
  print("yes" if m in yes_nums else "no")

ALDS1_5_B: Merge Sort

n = int(input())
S = [*map(int, input().split())]

def merge(A, left, mid, right):
  L = A[left:mid] + [float("inf")]
  R = A[mid:right] + [float("inf")]
  i, j = 0, 0
  for k in range(left, right):
    if L[i] <= R[j]:
      A[k] = L[i]
      i += 1
    else:
      A[k] = R[j]
      j += 1
  return right - left

def merge_sort(A, left, right):
  count = 0
  if right - left > 1:
    mid = (left + right) // 2
    count += merge_sort(A, left, mid) + merge_sort(A, mid, right)
    count += merge(A, left, mid, right)
  return count

count = merge_sort(S, 0, n)
print(*S)
print(count)

• マージソートは、ソートされた部分配列同士をマージして長さを倍にしていく安定なソート。マージ操作が$\Theta(n)$なので、最悪計算量が$\Theta(n\log n)$と最善。平均計算量も$\Theta(n\log n)$

• マージ操作では、∞の値を持つ番兵を両方の配列の最後に入れておくことで、一方が空になった場合の場合分けを省いている

• 比較操作のみに基づくソートでは、このオーダーは最悪計算量の下限。なぜなら、入力されるサイズ$n$の配列の順序関係は$\log (n!)=\Theta(n \log n)$ bitの情報量を持つ一方で、比較操作では1回あたり高々1bitの情報しか得られないため

ALDS1_5_C: Koch Curve

import cmath

def koch(n, p1, p2):
  if not n:
    return None
  s = (2*p1 + p2) / 3
  t = (p1 + 2*p2) / 3
  u = s + cmath.rect(1, cmath.pi/3) * (t - s)
  koch(n-1, p1, s)
  print(s.real, s.imag)
  koch(n-1, s, u)
  print(u.real, u.imag)
  koch(n-1, u, t)
  print(t.real, t.imag)
  koch(n-1, t, p2)

p1, p2 = 0j, 100+0j
print(p1.real, p1.imag)    
koch(int(input()), p1, p2)
print(p2.real, p2.imag)

• AOJではnumpyが使えないため行列演算はできないが、2次元の回転操作程度なら複素数で表現すればできる

ALDS1_5_D: The Number of Inversions

n = int(input())
A = [*map(int, input().split())]

def merge(A, left, mid, right):
  L = A[left:mid] + [float("inf")]
  R = A[mid:right] + [float("inf")]
  i, j, count = 0, 0, 0
  for k in range(left, right):
    if L[i] <= R[j]:
      A[k] = L[i]
      i += 1
    else:
      A[k] = R[j]
      j += 1
      count += mid - left - i
  return count

def merge_sort(A, left, right):
  count = 0
  if right - left > 1:
    mid = (left + right) // 2
    count += merge_sort(A, left, mid) + merge_sort(A, mid, right)
    count += merge(A, left, mid, right)
  return count

print(merge_sort(A, 0, n))

• ある配列Aを２つの連続した部分配列L, Rに分割するとき、Aの反転数は、L内の反転数と、R内の反転数と、LとRをまたぐ反転数の和になる。LとRをまたぐ反転数に関しては、LやRをソートしても値が変わらず、ソート済みの配列L, Rをまたぐ反転数なら、マージ操作中にRの各要素がLの元を何個飛び越すかをカウントしていけば分かる

トピック #6 ソート

ALDS1_6_A: Counting Sort

import itertools as it

n = int(input())
A = [*map(int, input().split())]

C = [0] * (max(A) + 1)
for a in A:
  C[a] += 1
C = [*it.accumulate(C)]

B = [None] * n
for a in reversed(A):
  C[a] -= 1
  B[C[a]] = a

print(*B)

• 計数ソートは、入力の取り得る値が整数でその範囲$k$が分かっている場合に使える安定なソート。取り得る各値に対してその値以下の要素数が何個あるかを計算することで、各入力をどれだけ右に配置すればよいかが分かる。最悪計算量は$O(n+k)$と高速で、比較操作のみに基づくソートでの下限$\Omega(n \log n)$を下回ることができる

ALDS1_6_B: Partition

n = int(input())
A = [*map(int, input().split())]

def partition(A, p, r):
  i = p - 1
  for j in range(p, r):
    if A[j] <= A[r]:
      i += 1
      A[i], A[j] = A[j], A[i]
  A[i+1], A[r] = A[r], A[i+1]
  return i+1

p = partition(A, 0, n-1)
print(*A[:p], f"[{A[p]}]", *A[p+1:])

• このように、配列をピボットと呼ばれる要素と、ピボット以下の要素全体と、ピボットより大きい要素全体の３つに分割する操作は$\Theta(n)$でできる

ALDS1_6_C: Quick Sort

n = int(input())
A = [[suit, int(num)] for suit, num in [input().split() for _ in range(n)]]

def partition(A, p, r):
  i = p - 1
  for j in range(p, r):
    if A[j][1] <= A[r][1]:
      i += 1
      A[i], A[j] = A[j], A[i]
  A[i+1], A[r] = A[r], A[i+1]
  return i+1

def quick_sort(A, p, r):
  if p < r:
    q = partition(A, p, r)
    quick_sort(A, p, q-1)
    quick_sort(A, q+1, r)

A_stable = sorted(A, key=lambda x:x[1])
quick_sort(A, 0, n-1)
print("Stable" if A == A_stable else "Not stable")
for x in A:
  print(*x)

• クイックソートは、配列を大きい部分配列と小さい部分配列に分割していく不安定なソート。最悪計算量は$\Theta(n^2)$と悪いけれど、平均計算量は$\Theta(n\log n)$と高速

ALDS1_6_D: Minimum Cost Sort

n = int(input())
w = [*map(int, input().split())]
ws = sorted(w)

cost = 0
for cyc_min in range(n):
    i = w.index(ws[cyc_min])
    cyc_l = 0
    while i > cyc_min:
        cyc_l += 1
        i_pre = w.index(ws[i])
        cost += w[i_pre]
        w[i], w[i_pre] = w[i_pre], w[i]
        i = i_pre
    cost += min(ws[cyc_min] * cyc_l, ws[cyc_min] * 2 + ws[0] * (cyc_l + 2))
print(cost)

• コストはサイクル毎に計算できる。cyc_minについてのループでは、cyc_minがどれかのサイクルの最小元であったときはそのサイクルのコストを計算し、どのサイクルの最小元でも無い場合は最初からiがcyc_minに一致して何もしないようにしてある。どのサイクルの最小元でも無い場合に最初からiがcyc_minに一致するようにするために、実際にサイクルを構成する互換(w[i],w[i_pre])を繰り返すことで既に訪れたサイクルに関してはwをwsに一致させている

• 各サイクルのコストは、そのサイクルの最小元を除いた各要素の重さと、運び屋のコストの和である。運び屋のコストは、運び屋としてそのサイクルの最小元ws[cyc_min]を使った場合のコストと、全体での最小元ws[0]をサイクルの外から借りてきて使って返した場合のコストのうち、小さい方を取る。cyc_lは各サイクルの長さ-1を表していて、長さ1のサイクルではcyc_lがゼロなのでコストはゼロになる