[SWEA] 1979 어디에 단어가 들어갈 수 있을까 - python

SW Expert Academy

SW Expert Academy

 

1. 가로부터 1이 있으면 cnt를 증가시킨다.

2. 한줄을 다 돌았을 때, cnt가 K와 일치하면 ans 를 증가시킨다.

-> 한줄을 다돌면 cnt 초기화

3. 전치행렬을 만들어서 세로도 똑같이 진행하면 끝

 

가로 / 세로 두번을 똑같은 코드를 반복하는 것이기 때문에

함수를 만들어서 2번 쓰는 것이 코드를 줄이는 방법

 

T = int(input())
for test_case in range(1, T+1):
    N, K = map(int,input().split())
    arr= [list(map(int,input().split())) for _ in range(N)]
    arr_t = list(map(list,zip(*arr)))
    def count(arr):
        ans = 0
        for i in range(N):
            cnt = 0
            for num in arr[i]:
                # 리스트 안에 요소가 1 일때
                if num == 1:
                    cnt +=1
                # 리스트 안에 요소가 0일 때
                else:
                    # 0만나면 cnt 초기화
                    # 만약 cnt == K면 ans +1하고 cnt 초기화
                    if cnt == K:
                        ans += 1
                    cnt = 0
            # 한줄을 다돌고 cnt == K 면 ans +1
            if cnt == K:
                ans += 1
        return ans
    ans = count(arr) + count(arr_t)
    print(f'#{test_case} {ans}')