|
| 1 | +# https://leetcode.com/problems/longest-balanced-subarray-ii/description/ |
| 2 | + |
| 3 | +''' |
| 4 | +You are given an integer array nums. |
| 5 | +A subarray is called balanced if the number of distinct even numbers in the subarray is equal to the number of distinct odd numbers. |
| 6 | +Return the length of the longest balanced subarray. |
| 7 | +
|
| 8 | +Example 1: |
| 9 | +Input: nums = [2,5,4,3] |
| 10 | +Output: 4 |
| 11 | +Explanation: |
| 12 | +The longest balanced subarray is [2, 5, 4, 3]. |
| 13 | +It has 2 distinct even numbers [2, 4] and 2 distinct odd numbers [5, 3]. Thus, the answer is 4. |
| 14 | +Example 2: |
| 15 | +Input: nums = [3,2,2,5,4] |
| 16 | +Output: 5 |
| 17 | +Explanation: |
| 18 | +The longest balanced subarray is [3, 2, 2, 5, 4]. |
| 19 | +It has 2 distinct even numbers [2, 4] and 2 distinct odd numbers [3, 5]. Thus, the answer is 5. |
| 20 | +Example 3: |
| 21 | +Input: nums = [1,2,3,2] |
| 22 | +Output: 3 |
| 23 | +Explanation: |
| 24 | +
|
| 25 | +The longest balanced subarray is [2, 3, 2]. |
| 26 | +It has 1 distinct even number [2] and 1 distinct odd number [3]. Thus, the answer is 3. |
| 27 | +
|
| 28 | +Constraints: |
| 29 | +1 <= nums.length <= 105 |
| 30 | +1 <= nums[i] <= 105 |
| 31 | +
|
| 32 | +Prefix Sum + Segment Tree |
| 33 | +''' |
| 34 | + |
| 35 | +class LazyTag: |
| 36 | + def __init__(self): |
| 37 | + self.to_add=0 |
| 38 | + |
| 39 | + def add(self,other): |
| 40 | + self.to_add+=other.to_add |
| 41 | + return self |
| 42 | + |
| 43 | + def has_tag(self): |
| 44 | + return self.to_add!=0 |
| 45 | + |
| 46 | + def clear(self): |
| 47 | + self.to_add=0 |
| 48 | + |
| 49 | +class SegmentTreeNode: |
| 50 | + def __init__(self): |
| 51 | + self.min_value=0 |
| 52 | + self.max_value=0 |
| 53 | + self.lazy_tag=LazyTag() |
| 54 | + |
| 55 | + |
| 56 | +class SegmentTree: |
| 57 | + def __init__(self,data): |
| 58 | + self.n=len(data) |
| 59 | + self.tree=[SegmentTreeNode() for _ in range(self.n*4+1)] |
| 60 | + self._build(data,1,self.n,1) |
| 61 | + |
| 62 | + def add(self,l,r,val): |
| 63 | + tag=LazyTag() |
| 64 | + tag.to_add=val |
| 65 | + self._update(l,r,tag,1,self.n,1) |
| 66 | + |
| 67 | + def find_last(self,start,val): |
| 68 | + if start>self.n: |
| 69 | + return -1 |
| 70 | + return self._find(start,self.n,val,1,self.n,1) |
| 71 | + |
| 72 | + def _apply_tag(self,i,tag): |
| 73 | + self.tree[i].min_value+=tag.to_add |
| 74 | + self.tree[i].max_value+=tag.to_add |
| 75 | + self.tree[i].lazy_tag.add(tag) |
| 76 | + |
| 77 | + def _pushdown(self,i): |
| 78 | + if self.tree[i].lazy_tag.has_tag(): |
| 79 | + tag=LazyTag() |
| 80 | + tag.to_add=self.tree[i].lazy_tag.to_add |
| 81 | + self._apply_tag(i<<1,tag) |
| 82 | + self._apply_tag((i<<1)|1,tag) |
| 83 | + self.tree[i].lazy_tag.clear() |
| 84 | + def _pushup(self,i): |
| 85 | + self.tree[i].min_value=min( |
| 86 | + self.tree[i<<1].min_value,self.tree[(i<<1)|1].min_value |
| 87 | + ) |
| 88 | + self.tree[i].max_value=max(self.tree[i<<1].max_value,self.tree[(i<<1)|1].max_value) |
| 89 | + |
| 90 | + def _build(self,data,l,r,i): |
| 91 | + if l==r: |
| 92 | + self.tree[i].min_value=data[l-1] |
| 93 | + self.tree[i].max_value=data[l-1] |
| 94 | + return |
| 95 | + mid=l+((r-l)>>1) |
| 96 | + self._build(data,l,mid,i<<1) |
| 97 | + self._build(data,mid+1,r,(i<<1)|1) |
| 98 | + self._pushup(i) |
| 99 | + |
| 100 | + def _update(self,target_l,target_r,tag,l,r,i): |
| 101 | + if target_l<=l and r<=target_r: |
| 102 | + self._apply_tag(i,tag) |
| 103 | + return |
| 104 | + self._pushdown(i) |
| 105 | + mid=l+((r-l)>>1) |
| 106 | + if target_l<=mid: |
| 107 | + self._update(target_l,target_r,tag,l,mid,i<<1) |
| 108 | + if target_r>mid: |
| 109 | + self._update(target_l,target_r,tag,mid+1,r,(i<<1)|1) |
| 110 | + self._pushup(i) |
| 111 | + |
| 112 | + def _find(self,target_l,target_r,val,l,r,i): |
| 113 | + if self.tree[i].min_value>val or self.tree[i].max_value<val: |
| 114 | + return -1 |
| 115 | + if l==r: |
| 116 | + return l |
| 117 | + self._pushdown(i) |
| 118 | + mid=l+((r-l)>>1) |
| 119 | + if target_r>=mid+1: |
| 120 | + res=self._find(target_l,target_r,val,mid+1,r,(i<<1)|1) |
| 121 | + if res!=-1: |
| 122 | + return res |
| 123 | + if l<=target_r and mid>=target_l: |
| 124 | + return self._find(target_l,target_r,val,l,mid,i<<1) |
| 125 | + return -1 |
| 126 | + |
| 127 | + |
| 128 | +class Solution: |
| 129 | + def longestBalanced(self, nums: List[int]) -> int: |
| 130 | + occurrences=defaultdict(deque) |
| 131 | + |
| 132 | + def sgn(x): |
| 133 | + return 1 if x%2==0 else -1 |
| 134 | + |
| 135 | + length=0 |
| 136 | + prefix_sum=[0]*len(nums) |
| 137 | + prefix_sum[0]=sgn(nums[0]) |
| 138 | + occurrences[nums[0]].append(1) |
| 139 | + |
| 140 | + for i in range(1,len(nums)): |
| 141 | + prefix_sum[i]=prefix_sum[i-1] |
| 142 | + occ=occurrences[nums[i]] |
| 143 | + if not occ: |
| 144 | + prefix_sum[i]+=sgn(nums[i]) |
| 145 | + occ.append(i+1) |
| 146 | + |
| 147 | + seg=SegmentTree(prefix_sum) |
| 148 | + for i in range(len(nums)): |
| 149 | + length=max(length,seg.find_last(i+length,0)-i) |
| 150 | + next_pos=len(nums)+1 |
| 151 | + occurrences[nums[i]].popleft() |
| 152 | + if occurrences[nums[i]]: |
| 153 | + next_pos=occurrences[nums[i]][0] |
| 154 | + seg.add(i+1,next_pos-1,-sgn(nums[i])) |
| 155 | + return length |
| 156 | + |
| 157 | + |
| 158 | + |
| 159 | +# Complexity Analysis |
| 160 | +# Let n be the length of nums. |
| 161 | +# Time complexity: O(nlogn). |
| 162 | +# Computing prefix sums and maintaining element positions require O(nlogn). Building the segment tree takes O(nlogn). Each iteration performs a constant number of segment tree operations and map accesses, each taking O(logn). Therefore, the total time complexity is O(nlogn). |
| 163 | + |
| 164 | +# Time complexity: O(n). |
| 165 | +# The segment tree, prefix sum array, and auxiliary data structures together require O(n) space. |
0 commit comments