Atcoder ABC241 振り返り

　Cがぱっと見難しそうでDがいけそうだったのでずっとDを解いていた．結局Dが解けずに2完．しかし，コンテスト終了後，＋3WAの＋20分くらいでCが解けたのでDにこだわらずCを解いておけば良かった……．Dは未知のアルゴリズムを理解する必要があったので後日しっかり勉強して解きます．

A - Digit Machine【AC】

　言われたとおりにやります．

a = [int(i) for i in input().split()]
a1 = a[0]
a2 = a[a1]
print(a[a2])

B - Pasta【AC】

　list.remove()は計算量多いので基本使いたくないですが，N<1000なので間に合います．for~else~と書くと，for文の中でbreakされずに処理を抜けたときのみ，else文の中の処理が行われます．便利なので最近よく使ってます．

N, M = map(int, input().split())
A = [int(i) for i in input().split()]
B = [int(i) for i in input().split()]

for b in B:
    if b in A:
        A.remove(b)
    else:
        print("No")
        break
else:
    print("Yes")

C - Connect 6【AC】

　マスを6つずつ見て.の数が2個以下ならYesを出力する．この方針がTLEしそうなのでDに行ってしまった．結局提出したら(結構ギリだけど)間に合ったので取りあえず1回提出してみることが大切だと学んだ．
　numpyを使ってマスを簡単に操作しています．横向きに見てできるか判定は素直にそのままやります．縦向きに見てできるか判定は，マスを転置して横向きに見て判定の処理をすれば，縦向きの判定になります．
　対角線に見てできるかの判定をします．numpyにはnp.diag()という関数があり簡単に対角成分を持ってくることが出来ます．さらに引数にkを設定することで開始地点を設定することが出来ます(ほんと便利)．これで/方向の対角線に見てできるか判定できます．最後に\方向の対角線でできるか判定する必要があります．これはマスを左右反転させて/方向に判定させれば良いと分かります．判定はnp.fliplr()で出来ます．
　これで全パターンについて判定が出来ました．バグを3WA分発生させるも下のコードでAC．コンテスト中にやっておけば良かったなぁ……

import numpy as np
N = int(input())
yoko_S = []
ok = False

def judge(arr):
    arr = list(arr)
    for i in arr:
        i = list(i)
        for j in range(N-5):
            if i[j:j+6].count(".") <= 2:
                return True
    return False

def diag_judge(arr):
    for i in range(-N, N):
        diag = list(np.diag(arr, k=i))
        if len(diag) < 6:
            continue
        for j in range(len(diag)-5):
            if diag[j:j+6].count(".") <= 2:
                return True
    return False

for _ in range(N):
    yoko_S.append(list(input()))
             
yoko_S = np.array(yoko_S)
tate_S = yoko_S.T
flip_S = np.fliplr(yoko_S)

if judge(yoko_S) or judge(tate_S) or diag_judge(yoko_S) or diag_judge(flip_S):
    print("Yes")
else:
    print("No")