Problem statement

https://binarysearch.com/problems/Longest-1s-After-One-Swap/

Solution

Variation of Leetcode 0487. Max Consecutive Ones II, we just need to check in the end that answer is no more than total number of ones.

Complexity

It is O(n) for time and O(1) for space.

Code

class Solution:
    def solve(self, A):
        ans, n, zeroes = 0, len(A), 0
        beg, end = 0, 0
        while end < n:
            if end < n and zeroes + (A[end] == "0")  <= 1:
                zeroes += (A[end] == "0")
                end += 1
                ans = max(ans, end - beg)
            else:
                zeroes -= (A[beg] == "0")
                beg += 1

        return min(ans, A.count("1"))