[SWEA] 오목판정 [PYTHON]

촤하하,,,
잘 풀려서 좋다 좋다 했더니
테스트 케이스 77개 이상을 맞추질 못하구나,,,

def check():

    # 가로
    for i in range(N):
        cnt = 0
        for j in range(N):
            if arr[i][j] == 'o':
                cnt += 1
            else:
                if cnt == 5:
                    return 'YES'
                cnt = 0
        if cnt == 5:
            return 'YES'

    #세로
    for i in range(N):
        cnt = 0
        for j in range(N):
            if arr[j][i] == 'o':
                cnt += 1
            else:
                if cnt == 5:
                    return 'YES'
                cnt = 0
        if cnt == 5:
            return 'YES'

    # 대각선 1
    cnt = 0
    for i in range(N):
        if arr[i][i] == 'o':
            cnt += 1
        else:
            if cnt == 5:
                return 'YES'
            cnt = 0
    if cnt == 5:
        return 'YES'


    # 대각선 2
    cnt = 0
    for i in range(N):
        if arr[i][N-1-i] == 'o':
            cnt += 1
        else:
            if cnt == 5:
                return 'YES'
            cnt = 0
    if cnt == 5:
        return 'YES'

    return 'NO'


for tc in range(1, int(input())+1):
    N = int(input())
    arr = [list(input()) for _ in range(N)]

    print('#{} {}'.format(tc, check()))
    

위의 코드에 어떤 문제로 정답을 맞추지 못했는지는 아직 미지수,

def check ():

    for i in range(N):

        for j in range(N-4):
            cnt = 0
            for k in range(5):
                if arr[i][j+k] == 'o':
                    cnt += 1
            if cnt == 5:
                return 'YES'

            cnt = 0
            for h in range(5):
                if arr[j+h][i] == 'o':
                    cnt += 1
            if cnt == 5:
                return 'YES'


    for k in range(N-4):
        for m in range(N-4):
            cnt = 0
            for z in range(5):
                if arr[k+z][m+z] == 'o':
                    cnt += 1
            if cnt == 5:
                return 'YES'

            cnt = 0
            for z in range(5):
                if arr[k+z][m+4-z] == 'o':
                    cnt += 1
            if cnt == 5:
                return 'YES'
    return 'NO'



for tc in range(1, int(input())+1):
    N = int(input())
    arr = [list(input()) for _ in range(N)]

    print('#{} {}'.format(tc, check()))

아마 대각선 부분의 테스트 케이스에서 걸린 것이 아닐까