Pythonで解く【初中級者が解くべき過去問精選 100 問】(015 - 017 全探索：順列全探索)

1. 目的

初中級者が解くべき過去問精選 100 問をPythonで解きます。
すべて解き終わるころに水色になっていることが目標です。

本記事は「015 - 017 全探索：順列全探索」です。

2. 総括

順列系の問題ではitertoolsをうまく使うと解ける場合が多いです。
permutaionsとcombinationとproductはよく使うので、この辺はいろいろサンプルデータで試しながら学びました。

3. 本編

015 - 017 全探索：順列全探索

015. AtCoder Beginner Contest 145 C - Average Length

回答

N = int(input())
towns = [tuple(map(int, input().split())) for _ in range(N)]

total_distance = 0
count = 0
for start in range(N-1):
    for end in range(start+1, N):
        start_x, start_y = towns[start]
        end_x, end_y = towns[end]

        total_distance += ((end_x - start_x)**2 + (end_y - start_y)**2)**0.5
        count += 1

# factor = math.factorial(N) とすると
# answer = total_distance * (factor * (N - 1) / count) / factor となり下記に式変形可能
answer = total_distance * ((N - 1) / count)
print(answer)

こちら計算方法がいろいろあります。
愚直にやるとpermutationですべての並べ方列挙して全てにおいて距離を求めてその平均をとります。
これでも通ると思いますが、もう少し計算量を落とした回答としてみました。
方針は、
1. すべての2点の組み合わせの距離を合計する
2. この合計に(N - 1) / 組み合わせの数を乗じたものが答え

です
上記2については、コードにコメントとして残した式変形から導きだしています。

---

016. AtCoder Beginner Contest 150 C - Count Order

回答

import itertools

N = int(input())
P = tuple(map(int, input().split()))
Q = tuple(map(int, input().split()))

count = 1
for p in itertools.permutations(range(1, N+1)):

    if p == P:
        count_P = count
    if p == Q:
        count_Q = count

    count += 1

ans = abs(count_P - count_Q)
print(ans)

pythonのitertools.permutationsが辞書順に列挙してくれることを信じて（！）愚直に書きます。
方針は、
1. 1~Nまでの順列をitertools.permutationsで列挙
2. 列挙するカウントを数えて置き、P、Qに一致したところでcount_Pとcount_Qを記録
3. 両者の差をとったものが答え

です。
itertoolsがない場合はもうちょっと考えなければいけないと思いますが、便利なので何も考えずに使用しました。

---

017. ALDS_13_A - 8 クイーン問題

回答

import itertools
import copy

def flag_queen_load(grid, r, c):
    for i in range(1, 8):
        if 0 <= c+i and c+i < 8:
            grid[r][c+i] = -1 #右
        if 0 <= c-i and c-i < 8:  
            grid[r][c-i] = -1 #左
        if 0 <= r+i and r+i < 8:  
            grid[r+i][c] = -1 #下
        if 0 <= r-i and r-i < 8:
            grid[r-i][c] = -1 #上
        if 0 <= r-i and r-i < 8 and 0 <= c+i and c+i < 8:
            if grid[r-i][c+i] == 9:
                continue
            grid[r-i][c+i] = -1 # 右上
        if 0 <= r+i and r+i < 8 and 0 <= c+i and c+i < 8:
            if grid[r+i][c+i] == 9:
                continue
            grid[r+i][c+i] = -1 # 右下
        if 0 <= r+i and r+i < 8 and 0 <= c-i and c-i < 8:  
            if grid[r+i][c-i]  == 9:
                continue
            grid[r+i][c-i] = -1 # 左下
        if 0 <= r-i and r-i < 8 and 0 <= c-i and c-i < 8:
            if grid[r-i][c-i] == 9:
                continue
            grid[r-i][c-i] = -1 # 左上

    grid[r][c] = 9 # 自分自身

    return grid


if __name__ == "__main__":

    k = int(input())
    grid = [[0] * 8 for _ in range(8)]
    for _ in range(k):
        r, c = map(int, input().split())
        grid = flag_queen_load(grid, r, c)

    # 0~7の番号を順列を全通り試す
    # 生成される順列について、添え字は行番号、要素は列番号として考える
    for perm in itertools.permutations(range(8)):
        copy_grid = copy.deepcopy(grid) # 毎回新しいgridを使う
        count = 0
        for row, col in enumerate(perm):
            if copy_grid[row][col] == -1:
                break
            if copy_grid[row][col] == 9:
                continue

            copy_grid[row][col] = 9
            copy_grid = flag_queen_load(copy_grid, row, col)
            count += 1

        if count == 8 - k: # 8-k個のクイーンをおけたらbreak
            break

    # 答え表示
    for row in range(8):
        for col in range(8):
            if copy_grid[row][col] == -1:
                print('.', end='')
            else:
                print('Q', end='')
        print()

もともとはNumpyで書いていたのですがAOJはNumpy使えないのでしょうか？エラーとなって提出できなかったので普通のリストで書き直しました。
方針は、
1. まずはチェス盤をgridとし、クイーンがある位置を9、クイーンが襲撃できる箇所を-1、それ以外の空白の個所を0とします
2. ある座標(r, c)にクイーンを置いた際に盤面を更新する関数flag_queen_load()を作成しておきます
3. あとはすべての置き方を試して、うまく置けるかを判定します
4. 全通り試すにあたってはitertools.permutations(range(8))で0~7の順列を列挙し、生成される順列について、添え字は行番号、要素は列番号として考えます
5. あとは、盤面が-1の場合は置けないのでbreakし、9の場合はOKなのでcontinueします
6. 新たにクイーンを8-k回置ける時が答え

です。

関数flag_queen_load()が汚いのでもうちょっときれいに書けばよかったです。

---