ラーメン、それは人類が到達した光

日々の気づきと思ったことをそのまま書き連ねるブログ

トップ > Python > Pythonで解く蟻本（1-6: 気軽にウォームアップ）

Pythonで解く蟻本（1-6: 気軽にウォームアップ）

Python Python-蟻本

自分の勉強の備忘録として書くので、ちゃんとした説明はないです。
僕の勉強の仕方はざっくりとこんな感じです。

  1. 蟻本の答案をまずはそのままPythonに変換
  2. 似たようなAtCoderの問題を自力で解いてみる
  3. 1日20のコードを目標とする

20コード、案外忘れるので自分のやった記録と共に記載していきます。
蟻本と似たような問題は、Qiitaの記事（AtCoder 版！蟻本 (初級編) - Qiita）を参考に探しました。

なお、持っている蟻本は以下の第2版です。

くじびき（P.8）

これは、ただ書き直すだけなので簡単です。
あえて言うなら、本の内容にある以下のような定数宣言がPythonに存在しないこと。

const int MAX_N = 50  //Pythonにはこんな宣言がない

なので PEP 8 に従って、定数として扱う変数（日本語おかしい気がするけど気にしない）は全部大文字（+アンダーバー）で表記しました。

www.python.org

このアルゴリズムは4重ループで計算量は O(n^4) なので、 n = 100 くらいになるとしんどくなります。（一般的に競プロだと 10e8 以下に計算量を抑えたいので）

MAX_N = 50  #Nの最大値を入力

n = int(input())
m = int(input())
k = list(map(int, input().split()))

f = False

for a in range(n):
  for b in range(n):
    for c in range(n):
      for d in range(n):
        if k[a] + k[b] + k[c] + k[d] == m:
          f = True

if f == True:
  print("Yes")
else:
  print("No")

三角形（蟻本 - P.21）

これも、シンプルな三重ループの問題です。
三重ループなので、n=500くらいまでは計算しきれる雰囲気っぽい。
※O(n^3) の計算量

n = int(input())
a = list(map(int, input().split()))

ans = 0
# forループを、ダブらないように一つずつズラしている
for i in range(n - 2):
  for j in range(i + 1, n - 1):
    for k in range(j + 1, n):
      len = a[i] + a[j] + a[k]
      ma = max(a[i], max(a[j], a[k]))
      rest = len - ma

      if ma < rest:
        ans = max(ans, len)

print(ans)

ARC 004 A 2点間距離の最大値

A - 2点間距離の最大値 ( The longest distance )

特に深く考えずにできる問題です。
二重ループでも、N <= 100 なので計算間に合います。

個人的に、平方根の計算に関しては math.sqrt() を使うよりはデフォの演算子 ** を使って1/2乗として計算する方が好みです。
n乗根の計算になっても焦らず同じ方法で計算できるので。

n = int(input())
x = []
y = []

for i in range(n):
  xin, yin = map(int, input().split())
  x.append(xin)
  y.append(yin)

max_d = 0

for i in range(n):
  for j in range(i + 1, n):
    d = ((x[i]-x[j])**2 + (y[i]-y[j])**2) ** (1/2)
    max_d = max(max_d, d)

print(max_d)

ABC 051 B Sum of Three Integers

B - Sum of Three Integers

これは少しひねりが必要です。
3重ループを普通に実装すると、 n = 500 程度でTLEしてしまいます。
S - (X + Y) が 0~Kの範囲にあるなら組み合わせとしてカウントする、としてしまえば良いので、以下のコードになります。

K, S = map(int, input().split())
count = 0

for i in range(K + 1):
  for j in range(K + 1):
    if S - (i + j) <= K and S - (i + j) >= 0:
      count += 1

print(count)

ABC 085 C Otoshidama

C - Otoshidama

これも、一つ上の問題と同じ発想で解けますね。

N, Y = map(int, input().split())

count = 0
x, y, z = -1, -1, -1

for i in range(N + 1):
  for j in range(N + 1 - i):
    if Y - (i*10000 + j*5000) == (N - i - j)*1000:
      x, y, z = i, j, N - i - j
      break
  else:
    continue
  break

print(x, y, z)

まとめ

単純なループ処理の解き方が思いついた場合は、まずはループ処理を書いてみた方がイメージしやすい。
その上で、計算量の概算からループをいくつ以下に抑えるべきか考えて、内側のループを減らせないか検討すると良さそう。

Ants（蟻本 - P.23）

この問題は凄いですね。。
本の解説が無限にわかりやすいのでコードだけ載せます。

L = int(input())
n = int(input())
x = list(map(int, input().split()))

minT = 0
for i in range(n):
  minT = max(minT, min(x[i], L - x[i]))

maxT = 0
for i in range(n):
  maxT = max(maxT, max(x[i], L - x[i]))

print(minT, maxT)

AGC 013 C Ants on a Circle

C - Ants on a Circle

めちゃくちゃ難しいみたいです。
実際、めちゃくちゃ難しかったので解説を読みました。これは致し方なし。

n,l,t=map(int,input().split())
arr=[list(map(int,input().split())) for _ in range(n)]
ans=[]
cnt=0

for i in range(n):
  tmp=arr[i][0]  #蟻iの最初の座標
  # 衝突を考慮しなかった場合の各蟻のt秒後の座標をtmpに入れる
  # cntには円を周回した回数の合計値を入れる
  if arr[i][1]==1:
    tmp+=t
    tmp%=l
    cnt+=(arr[i][0]+t)//l
  else:
    tmp-=t
    tmp%=l
    cnt+=(arr[i][0]-t)//l
  ans.append(tmp)
ans=sorted(ans) #ansに格納したtmp（座標）を昇順にソート
cnt%=n #0番目目の蟻の衝突回数
tmp=[0]*n
for i in range(n):
  #初期位置0番目の蟻の最終位置は、原点から見て衝突数番目の位置にいるはず
  #各蟻の位置関係は、衝突を繰り返すので変わらず時計回りに位置するはず
  tmp[i]=ans[(cnt+i)%n]
ans=tmp
for pos in ans:
  print(pos)

まとめ

難しい問題に遭遇したら、まずは一度簡単に捉え直す！

ハードルの上がった「くじびき」（蟻本 - P.25）

最初のくじびきの n <= 50 だったのが、 n <= 1000になったバージョンです。
冒頭の4重ループだと、TLEしてしまうのでループを2つ削りたいです。

三角形の問題の例題で解いたAtCoderの問題と同じ発想で、まずは3重ループにします。

MAX_N = 50  #Nの最大値を入力

n = int(input())
m = int(input())
k = list(map(int, input().split()))

f = False

for a in range(n):
  for b in range(n):
    for c in range(n):
      d = m - (k[a]+k[b]+k[c])
      if d in k:  #ここの計算量がO(n)
        f = True
        break
  else:
    continue
  break

if f == True:
  print("Yes")
else:
  print("No")

ただし、最後の内側のループで使っている d in k という部分の計算量は O(n) です。
つまり、結局アルゴリズム全体の計算量は4重ループと変わりません。

今回検索したいのは数字なので二分探索を行うことで検索の計算量を O(n) → O(log n) に落とすことができます。
事前のソートにも O(n) の計算量がかかりますがループの外側でソート可能なので、無視できます。
ということで、修正すると以下のコードになります。
※計算量は O(n^3 logn)

def binary_search(x: int, k: list):
  left = 0
  right = n

  while right - left >= 1:
    mid = (left + right) // 2
    if k[mid] == x:
      return True
    elif k[mid] < x:
      left = mid + 1
    else:
      right = mid

  return False

n = int(input())
m = int(input())
k = list(map(int, input().split()))
k.sort()

f = False

for a in range(n):
  for b in range(n):
    for c in range(n):
      d = m - (k[a]+k[b]+k[c])
      if binary_search(d, k):
        f = True
        break
  else:
    continue
  break

if f == True:
  print("Yes")
else:
  print("No")

これでもまだ計算量が足りないので、更に同じ発想で3つ目のループの内側も二分探索とします。
（k[c]も二分探索で実装するイメージ）

ちなみに、先ほどのコードでは蟻本に倣って自前で二分探索を実装しましたが、pythonにはbisectという公式ライブラリがあるので今回はそれを使う形の修正も入れます。

import bisect

MAX_N = 1000
kk = []

n = int(input())
m = int(input())
k = list(map(int, input().split()))
k.sort()

f = False

for c in range(n):
  for d in range(n):
    kk.append(k[c] + k[d])

kk.sort()
for a in range(n):
  for b in range(n):
    cd = m - (k[a]+k[b])
    if bisect.bisect(kk, cd) > 0:
      f = True
      break
  else:
    continue
  break

if f == True:
  print("Yes")
else:
  print("No")

これで、めでたく計算量 O(n^2 logn) に落とすことができました。

AOJ 0529

Darts

上記のハードルの上がったくじびきと同じ要領でやってみました。
ただし、これではTLEしてしまいます。

import bisect

while True:
  max_S = 0
  p = [0]
  pp = []
  n, m = map(int, input().split())
  if n == 0 and m == 0:
    break

  for _ in range(n):
    p.append(int(input()))

  for c in range(n + 1):
    for d in range(c, n + 1):
      pp.append(p[c] + p[d])

  pp.sort()

  for a in range(n + 1):
    for b in range(a, n + 1):
      S_tmp = p[a] + p[b]
      l = bisect.bisect(pp, m - S_tmp)
      if l > 1:
        max_S = max(max_S, S_tmp + pp[l - 1])
      if max_S == m:
        break
    else:
      continue
    break

  print(max_S)

ここで、ちゃんと考えてみると最初の二重ループと二回目の二重ループは実はほぼ同じことをしています。
なので、上記のコードのppを一度計算させておいて、ppについて1度ループを回せば良いことがわかります。

更に、そもそもppがmを超える場合は考慮しなくて良さそうです。
ということは、こうなります。（結構、ギリギリの秒数になってしまいますがこれでAcceptされました。）
※計算量はたぶん、O(n^2 log n) です。

import bisect

while True:
  max_S = 0
  p = [0]
  pp = []

  n, m = map(int, input().split())
  if n == 0 and m == 0:
    break

  for _ in range(n):
    p.append(int(input()))

  for c in range(n + 1):
    for d in range(c, n + 1):
      pp.append(p[c] + p[d])

  pp.sort()

  # ppの中でm以下の最大数の位置 + 1
  t = bisect.bisect_left(pp, m)

  if t < len(pp): #m以上の数字が1つでもある場合
    if pp[t] == m: #mがあるならそのまま投げなきゃ良い
      print(m)
      continue
    pp = pp[:t] #m以下の数字セットに絞る
  max_S = 0

  for i in range(len(pp)):
    t = bisect.bisect_left(pp, m - pp[i])
    if pp[i] + pp[t - 1] > max_S:
      max_S = pp[i] + pp[t - 1]
  print(max_S)

まとめ

  1. forループの計算結果、使いまわせるなら使い回す
  2. 内側のループは頑張って引き算で表現できないか考える
  3. 二分探索で計算量落として頑張る

yassie-tp-0243

関連記事

コメントを書く