one line of code at a time

[leetcode] 1004. Max Consecutive Ones III 파이썬 코드 본문

leetcode

[leetcode] 1004. Max Consecutive Ones III 파이썬 코드

oloc 2024. 8. 13. 09:58

https://leetcode.com/problems/max-consecutive-ones-iii

 

0과 1로 이루어진 배열 nums가 있다. 0을 1로 바꿀 수 있는 횟수인 k가 주어진다. 최대 k만큼 0을 1로 바꿨을 때 1이 연속적으로 나오는 subarray의 최대 길이를 구하는 문제다. 

 

nums = [1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0] 이고, k = 2일 때 5번째 인덱스와 맨 마지막 인덱스에 있는 0을 1로 바꿔서 길이가 6인 [1, 1, 1, 1, 1, 1] 배열을 만들 수 있다. 그러므로 정답은 6이다. 

 

고민하다가 안 풀려서 다른 사람 코드를 참고했다. 

 

left pointer와 right pointer를 이용한다. left pointer는 0으로 설정해둔다. 

for 문을 돌리면서 right pointer를 하나씩 증가시켜준다. 

 

right pointer가 3번째 인덱스에 왔을 때 nums[rp] == 0이다. 

0의 개수도 카운트 해준다. num_zero를 1 증가시켜준다. 

num_zero는 최대 k개까지 가능하다. 

 

num_zero가 k보다 커지면, left pointer를 조정해서 (1씩 증가시켜서) 다시 valid하게 sliding window를 만들어야 한다. 

 

class Solution(object):
    def longestOnes(self, nums, k):
        num_zero = 0
        max_length = 0
        lp = 0

        for rp in range(len(nums)):
            if nums[rp] == 0:
                num_zero += 1

            while num_zero > k: # window is invalid
                if nums[lp] == 0:
                    num_zero -= 1
                lp += 1

            # window is valid
            max_length = max(max_length, rp - lp + 1)
        
        return max_length

 

 

참고

https://www.youtube.com/watch?v=HsGKI02yw6M