[코드트리 파이썬] 1이 K개 이상 존재하는 부분 수열

문제

---

풀이

투 포인터를 이용하였다.

1이 K개 이상 존재하는 가장 짧은 수열의 길이를 구하여야 한다.

 

j = 0
oneCnt = 0

먼저 투포인터의 두번째 포인터인 j와 1의 개수를 세어줄 oneCnt를 선언하였다.

 

1의 개수를 세어주기 위하여 다음과 같은 방법을 사용하였다.

먼저 첫 번째 포인터인 i를 순회시킨다.

이 때 i가 1을 만난다면 두 번째 포인터를 순회시킨다.

for i in range(n):
    if arr[i] == 1:
        while j < n and oneCnt != k:
            if arr[j] == 1:
                oneCnt += 1
            j += 1

두 번째 포인터는 끝까지 도달했거나 1을 k만큼 카운트할 때까지 순회한다.

순회 중 1을 만나면 oneCnt값을 1 증가시킨다.

 

 

if oneCnt != k and res == 4e9:
            res = -1
            break

만약 res값이 초기값과 같고 1의 개수가 k만큼 도달하지 못했으면 

처음 주어진 배열에 조건을 만족하는 부분 수열이 없다는 것이므로 

res에 -1을 넣고 투포인터를 종료한다.

if oneCnt == k:
	res = min(res, j - i)
oneCnt -= 1

 

만약 1의 개수가 k값에 도달했으면 res값을 갱신해준다.

또한 다음 수열을 찾기 위해 i값을 이동시켜야 하고, 현재 i값이 가리키고 있는건 1이므로 1의 개수를 하나 줄여주고 i를 이동시킨다.

 

---

코드

n, k = map(int, input().split())

arr = list(map(int, input().split()))


j = 0
oneCnt = 0
res = 4e9
for i in range(n):
    if arr[i] == 1:
        while j < n and oneCnt != k:
            if arr[j] == 1:
                oneCnt += 1
            j += 1
        
        if oneCnt != k and res == 4e9:
            res = -1
            break
        
    

        if oneCnt == k:
            res = min(res, j - i)
        oneCnt -= 1

print(res)