見出し画像
Photo by kimusubi

Pythonの便利な関数4つ+αを一題で復習 (ABC326 D - ABC Puzzle)

こじま

はじめに

競技プログラミングの話です。Pythonには便利な関数が標準で備えられています。上手に使えると、できることが広がってとても楽しいです。複雑な処理をほんの僅かなコードで書けることもあります。コーディング速度もきっと向上するでしょう。

しかしながら、関数を使う機会が少ないままだとなかなか覚えられないものです。使いこなすためには、くりかえし触れて慣れる必要があるでしょう。

この記事では、次の4つの関数を学びます。

  • itertools.product : 複数のリストの直積を作る

  • any : 1つでも真があればTrueを返す

  • enumerate : リストのインデックスと要素をまとめて取得する

  • map : リストの要素に関数・処理を適用する

とても身近な関数が1つ含まれているかもしれませんが……これら4つの関数を、一題で復習できる問題があります。パナソニックグループ プログラミングコンテスト2023（AtCoder Beginner Contest 326）のD問題です。この問題の解説コードを、4つのパートに分けて読み解いていきます。

パートごとに、便利な関数を使わずに同じ処理ができるコードを紹介します。基本的な関数を用いて書き直すことで、関数の意味を直感的に理解できると思います。より詳細を知りたい場合は、文末の『参照』をご活用ください。

AtCoder灰〜緑向けの内容です。ただし、リスト内包表記とジェネレータ式の基礎知識を前提としています。必要に応じて、以下のリンクなどをご参照ください。


問題 

$${N \times N}$$ の空のマス目があります。文字 A,B,C から1文字を選んで、それぞれのマス目に書き込むことができます。文字を書き込まずに空のままでもOKです。

このとき、以下の条件を満たすことができるか判定して、可能な場合は書き込み方を一つ示してください。

  • 各行・各列に A,B,C がちょうど1個含まれる。

  • 各行の左端の文字を縦読みすると、文字列$${R}$$と一致する。

  • 各列の上端の文字を横読みすると、文字列$${C}$$と一致する。

制約

  • $${3 \le N \le 5}$$

  • $${R,C}$$は A,B,C からなる長さ$${N}$$の文字列

解説コードを読み解く

kyopro_friendsさんのユーザ解説とコードを全面的に参考にしています。
最初に、コード全体を一度引用します。

import itertools
N=int(input())
R=input()
C=input()

def head(s):
  for c in s:
    if c!=".":
      return c

for a,b,c in itertools.product(itertools.permutations(range(N)),repeat=3):
  if any(ai==bi or bi==ci or ci==ai for ai,bi,ci in zip(a,b,c)): continue
  S=[["."]*N for _ in range(N)]
  for i,j in enumerate(a): S[i][j]="A"
  for i,j in enumerate(b): S[i][j]="B"
  for i,j in enumerate(c): S[i][j]="C"
  if "".join(map(head,S))==R and "".join(map(head,zip(*S)))==C:
    print("Yes")
    print(*map(lambda s:"".join(s),S),sep="\n")
    exit()
print("No")

このコードを4つのパートに分けて読んでいきます。

(1) itertools.productで全探索

「各行・各列に A,B,C がちょうど1個含まれる」という条件を使って順列全探索をします。

import itertools

for a,b,c in itertools.product(itertools.permutations(range(N)),repeat=3):

で、全パターンを回すことができます。ここで、順列 $${a}$$ の$${i}$$番目の値$${j}$$は、$${i}$$行$${j}$$列に文字 A を配置することに対応しています。$${b, c}$$ についても同様です。

この問題ではマス目に入る文字は3種類で固定されているので、productの繰り返し回数も常に3回です。そのため、forループを3つ重ねればproductの代用ができます。

from itertools import permutations

# 順列全探索 productなしで
for a in permutations(range(N)) :
    for b in permutations(range(N)) :
        for c in permutations(range(N)) :

（参考：itertools.permutationsを重ねて使う方法もあります。ループを回す回数がすこし減ります。順列$${a, b, c}$$はそれぞれ異なる必要があることを利用しています。）

from itertools import permutations

# 順列全探索 productなしで
P = permutations(range(N))
for a, b, c in permutations(P,3) :

(2) anyで枝刈り

(1)のproductによる順列全探索には、同じ位置に別の文字が重複する場合が含まれています。各マス目には一文字しか置けないので、重複を枝刈りします。

  if any(ai==bi or bi==ci or ci==ai for ai,bi,ci in zip(a,b,c)): continue

anyの引数には、()を省略したジェネレータ式が使われています。見た目もすっきりしていますし、リスト内包表記よりも高速になるそうです。

anyを使わないとしたら、一例として以下のコードが考えられます。

# 枝刈り anyなしで
    flag = False
    for ai,bi,ci in zip(a,b,c) :
        if ai==bi or bi==ci or ci==ai :
            flag = True
            break
    if flag : continue

(3) enumerateで順列から配置を作る

さて、順列 $${a,b,c}$$ から具体的な文字の配置 $${S}$$ を作ります。

    S = [["."]*N for _ in range(N)]
    for i,j in enumerate(a) : S[i][j] = "A"
    for i,j in enumerate(b) : S[i][j] = "B"
    for i,j in enumerate(c) : S[i][j] = "C"

enumerateが便利です。繰り返しになりますが、順列$${a}$$の$${i}$$番目の値$${j}$$は、$${i}$$行$${j}$$列に文字 A を配置することに対応しています。

ここはfor文と添字でもシンプルに書けます。

# 配列を作る enumerateなしで
    S = [["."]*N for _ in range(N)]
    for i in range(N) :
        S[i][a[i]] = "A"
        S[i][b[i]] = "B"
        S[i][c[i]] = "C"

(4) mapで縦読み・横読みの判定と出力

$${S}$$ が縦読みと横読みの条件を満たしているかを判定して、合致した場合に$${S}$$を出力して終了します。

まず、関数 head を定義します。

def head(s):
  for c in s:
    if c!=".":
      return c

リストなどを受け取って、空マス "." でない最初の文字を返す関数です。この関数をmapと組み合わせることで、判定を行います。

  if "".join(map(head,S))==R and "".join(map(head,zip(*S)))==C:
    print("Yes")
    print(*map(lambda s:"".join(s),S),sep="\n")
    exit()

mapの第一引数には、組み込み関数だけでなく、自分で定義した関数やラムダ式を入れることができます。

mapの代わりとして、リスト内包表記でも同じ処理ができます。

# 判定と出力 mapなしで
    r = "".join([head(s) for s in S])
    c = "".join([head(s) for s in zip(*S)])
    if r == R and c == C :
        print("Yes")
        ans = ["".join(s) for s in S]
        print(*ans,sep="\n")
        exit()

振り返り・まとめ

mapを入力以外で使ったことがなかったので、解説コード中の

"".join(map(head,S))==R

の部分がとても印象に残った問題でした。調べるうちに、リスト内包表記とジェネレータ式の働きや、便利な関数の使い方についても学ぶことができたので、記事としてまとめてみました。

この記事がよい復習の機会となれば幸いです。

参照

Pythonの文法はnkmkさんのサイトを参考にしています。いつもお世話になっています。感謝。

また、これら4つの関数を使わずに、代替方法でまとめたACコードはこちらです。

以上です。最後まで読んでいただき、ありがとうございました。

この記事が気に入ったらサポートをしてみませんか？