Longest Balanced Substring After One Swap | Longest Balanced Substring After One Swap

class Solution {
    int findLongest(string& s){
        int n=s.size();
        int ones=0,zeroes=0;
        for(int i=0;i<n;++i){
            if(s[i]=='1')   ones++;
            else            zeroes++;
        }
        int longest = 0;
        int balance = 0;
        unordered_map<int,int> first_idx;
        first_idx[0] = -1;
        for(int i=0;i<n;++i){
            if(s[i]=='1'){
                ones--;
                balance++;
            }else{
                zeroes--;
                balance--;
            }
            if(balance==0){// Prefix[i] is the longest substring
                longest = i+1;
                continue;
            }

            // -2 Case: Increasing Gradient (Need a 0 from right)
            if(first_idx.count(balance-2) and zeroes>0)
                longest = max(longest,i-first_idx[balance-2]);
            
            // +2 Case: Decreasing Gradient (Need a 1 from right)
            if(first_idx.count(balance+2) and ones>0)
                longest = max(longest,i-first_idx[balance+2]);
            
            // Insert new entry: If Unique
            if(!first_idx.count(balance))
                first_idx[balance] = i;
        }
        return longest;
    }
public:
    int longestBalanced(string s) {
        int longest = findLongest(s);
        reverse(s.begin(),s.end());
        longest = max(longest, findLongest(s));
        return longest;
    }
};
/*
//JAVA
import java.util.*;

class Solution {
    private int findLongest(String s) {
        int n = s.length();
        int ones = 0, zeroes = 0;
        
        // Initial count of ones and zeroes
        for (char c : s.toCharArray()) {
            if (c == '1') ones++;
            else zeroes++;
        }

        int longest = 0;
        int balance = 0;
        Map<Integer, Integer> firstIdx = new HashMap<>();
        firstIdx.put(0, -1);

        for (int i = 0; i < n; i++) {
            if (s.charAt(i) == '1') {
                ones--;
                balance++;
            } else {
                zeroes--;
                balance--;
            }

            if (balance == 0) {
                // Prefix ending at i is balanced
                longest = i + 1;
            } else {
                // -2 Case: Check if balance-2 existed and we still have zeroes left
                if (firstIdx.containsKey(balance - 2) && zeroes > 0) {
                    longest = Math.max(longest, i - firstIdx.get(balance - 2));
                }
                
                // +2 Case: Check if balance+2 existed and we still have ones left
                if (firstIdx.containsKey(balance + 2) && ones > 0) {
                    longest = Math.max(longest, i - firstIdx.get(balance + 2));
                }
            }

            // Store the first occurrence of this balance
            if (!firstIdx.containsKey(balance)) {
                firstIdx.put(balance, i);
            }
        }
        return longest;
    }

    public int longestBalanced(String s) {
        int res1 = findLongest(s);
        String reversed = new StringBuilder(s).reverse().toString();
        int res2 = findLongest(reversed);
        return Math.max(res1, res2);
    }
}

#Python
class Solution:
    def findLongest(self, s: str) -> int:
        n = len(s)
        ones = s.count('1')
        zeroes = n - ones
        
        longest = 0
        balance = 0
        # Stores first occurrence of each balance
        first_idx = {0: -1}
        
        for i, char in enumerate(s):
            if char == '1':
                ones -= 1
                balance += 1
            else:
                zeroes -= 1
                balance -= 1
            
            if balance == 0:
                longest = i + 1
            else:
                # -2 Case: Increasing Gradient
                if (balance - 2) in first_idx and zeroes > 0:
                    longest = max(longest, i - first_idx[balance - 2])
                
                # +2 Case: Decreasing Gradient
                if (balance + 2) in first_idx and ones > 0:
                    longest = max(longest, i - first_idx[balance + 2])
            
            # Insert only if the balance hasn't been seen before
            if balance not in first_idx:
                first_idx[balance] = i
                
        return longest

    def longestBalanced(self, s: str) -> int:
        # Check original string and its reverse
        return max(self.findLongest(s), self.findLongest(s[::-1]))
*/