きゃねろぐ

とある「おひさま」の徒然日記

Python初心者のためのABC201

A問題

atcoder.jp

 A_1\le A_2\le A_3 となるように並べ替えた後,等差数列になっているかを確認すれば良い. もちろん, A_1\ge A_2\ge A_3 となるように並べ替えても構わない.

A = list(map(int,input().split()))

A.sort()

if A[2]-A[1]==A[1]-A[0]:
  print('Yes')
else:
  print('No')

別解を示す.条件式を変形すると,

 A_3-A_2=A_2-A_1
 A_1+A_3=2A_2
 3(A_1+A_3)=2(A_1+A_2+A_3)


となる. A_1 A_3 A の中の最小値と最大値であることを利用して,以下のように書ける.

A = list(map(int,input().split()))

if 3*(min(A)+max(A))==2*sum(A):
  print('Yes')
else:
  print('No')

B問題

atcoder.jp

Pythonのソートsortにはkeyという引数があり,並び替えるときの基準を指定することができる.

# 7で割った余りを出力する関数
def f(x):
    return x % 7

a=[2, 5, 8, 14]
a.sort(key = f)   # keyで指定した関数の値に応じて並べ替える

print(a)
> [14, 8, 2, 5]   # 7で割った余りが小さい順にソートされる

今回の場合,各要素が  [S,T](  S は山の名前, T は山の高さ)となっており, T の値の降順に並び替える必要がある. 普通にsortを用いると昇順にソートされるが,引数にreverse = Trueと指定することで降順にソートされる.

N = int(input())
mountain = []

for i in range(N):
  S,T = input().split()
  T = int(T)
  mountain.append([S,T])

# 各要素X=[S,T]が与えられたときにT=X[1]の値に応じてソートしたい 
def f(X):
  return X[1]

mountain.sort(reverse = True, key = f)
print(mountain[1][0])

別解として,山の名前と高さでリストを分けておいて,高さだけソートして2番目に高い山の高さを計算した後に山の名前を調べても良い.

N = int(input())
name = {}
height = []

for i in range(N):
  S,T = input().split()
  T = int(T)
  name[T] = S   # 高さTの山の名前はSである
  height.append(T)
  
height.sort(reverse = True)
h = height[1]   # 2番目に高い山の高さ
print(name[h])

C問題

あらゆる問題を解くことができる最強の方法は,有り得るパターンをすべて試す全探索である.唯一の欠点は,パターンの数に比例した多大な計算コストを要することである.AtCoderの問題を見たら,まずは「計算量的に全探索できるかどうか?」を考える習慣をつけておくと,全探索で解ける問題が出たときに解法を考える時間を省いて瞬時に実装へと進むことができる.

今回の場合は調べるべき暗証番号の組み合わせは  10^4=10000 通りと少ないため,全探索が可能である.暗証番号全てを生成する方法として,最も単純なものとしてfor pin in range(10000)が考えられるが,

  • 整数から○桁目の数を取り出す操作に手間がかかる(一旦文字列に変換する,商と余りを用いる,など)
  • 0埋めの操作が必要(例えば,pin = 123は暗証番号「0123」として扱う必要がある)

といった点に注意する必要がある.今回は,代わりに テクニック集 でも紹介したitertools.productを使ってみる.

from itertools import product

# 0,1,...,9から4回要素を取り出す全ての組み合わせをタプルで出力する
for pin in product(range(10), repeat = 4):
  print(pin)

> (0, 0, 0, 0)
> (0, 0, 0, 1)
> (0, 0, 0, 2)
> (0, 0, 0, 3)
> (0, 0, 0, 4)
       ︙

なお,タプル()はリスト[]と似ているが,いくつか違いがある.

  • 要素の変更ができない(当然appendなども使えない)
  • dictkeyとして指定できる

暗証番号は全て生成できるようになったので,後は各暗証番号が条件をみたすかどうかをチェックすれば良い.ポイントは

  1. oの数字は全て暗証番号に現れる
  2. 暗証番号に現れる数字は全てxでない

の2点である.oの数字,xの数字をリストなどで管理して調べていけば良い.

from itertools import product

S = input()
maru = []
batsu = []

for i in range(10):
  if S[i]=='o':
    maru.append(i)
  if S[i]=='x':
    batsu.append(i)

ans = 0

# 0,1,...,9から4回要素を取り出す全ての組み合わせをタプルで出力する
for pin in product(range(10), repeat = 4):
  
  ok = True   # 条件をみたすかどうか
  
  # 条件1
  for i in maru:
    if i not in pin:
      ok = False
      
  # 条件2
  for i in pin:
    if i in batsu:
      ok = False
      
  if ok:
    ans += 1
    
print(ans)

別解として,oの個数とxの個数から組み合わせ論的に答えを導出できるが,今回の場合は全探索のコードを書いてしまう方が早いだろう.