Day 25 - Check If All 1's Are at Least Length K Places Away

#leetcode #python #challenge

The Problem

Given an array nums of 0s and 1s and an integer k, return True if all 1's are at least k places away from each other, otherwise return False.

Example 1:

Input: nums = [1,0,0,0,1,0,0,1], k = 2
Output: true
Explanation: Each of the 1s are at least 2 places away from each other.

Example 2:

Input: nums = [1,0,0,1,0,1], k = 2
Output: false
Explanation: The second 1 and third 1 are only one apart from each other.

Example 3:

Input: nums = [1,1,1,1,1], k = 0
Output: true

Example 4:

Input: nums = [0,1,0,1], k = 1
Output: true

Constraints:

1 <= nums.length <= 10^5
0 <= k <= nums.length
nums[i] is 0 or 1

Tests

import pytest
from .Day25_CheckIfAll1sAreAtLeastLengthKPlacesAway import Solution

s = Solution()


@pytest.mark.parametrize(
    "nums,k,expected",
    [
        ([1,0,0,0,1,0,0,1], 2, True),
        ([1,0,0,1,0,1], 2, False),
        ([1,1,1,1,1], 0, True),
        ([0,1,0,1], 1, True)
    ],
)
def test_k_length_apart(nums, k, expected):
    assert s.kLengthApart(nums, k) == expected

Solution

from typing import List


class Solution:
    def kLengthApart(self, nums: List[int], k: int) -> bool:
        current_dist = k
        for n in nums:
            if n == 1:
                if current_dist < k:
                    return False
                current_dist = 0
            else:
                current_dist += 1

        return True