Finding balance in a sea of numbers is a classic algorithmic challenge that tests your ability to track state changes over time. In this problem, we move beyond simple counting and dive into the world of distinct elements, requiring a clever mix of data structures to maintain efficiency.
Problem Summary
You're given:
- An array of integers called
nums.
Your goal:
- Find the length of the longest subarray where the count of distinct even numbers equals the count of distinct odd numbers.
Intuition
The "distinct" requirement is the main hurdle here. If we only needed to count occurrences, a simple prefix sum would suffice. However, because we only care about the first time a number appears in a subarray, the contribution of a number depends on the subarray's boundaries.
We can represent the "balance" by treating distinct odd numbers as and distinct even numbers as . A subarray is balanced if the sum of these contributions is 0.
As we slide our starting index to the right, a number nums[i-1] that was previously "distinct" for subarrays starting at is no longer included. However, its next occurrence later in the array now becomes the "first" appearance for subarrays starting at . To handle these range updates efficiently, we use a Segment Tree with Lazy Propagation. This allows us to update the balance values across multiple potential end positions simultaneously and quickly find the rightmost index where the balance hits zero.
Walkthrough: Understanding the Examples
Consider Example 3: nums = [1, 2, 3, 2]
- Initial State (Subarrays starting at index 0):
- Index 0 (1): Odd, count = 1. Balance = 1.
- Index 1 (2): Even, count = 1. Balance = . (Length 2:
[1, 2]is balanced). - Index 2 (3): Odd, count = 2. Balance = .
- Index 3 (2): Even, but 2 is already seen. Balance stays 1.
Max length so far: 2.
Shift to start at index 1:
The 1 at index 0 is removed. Its contribution (+1) is subtracted from all subsequent balances.
New balances for subarrays starting at index 1:
Index 1 (2): Even, balance = -1.
Index 2 (3): Odd, balance = . (Length 2:
[2, 3]is balanced).Index 3 (2): Since the 2 at index 1 is now "outside", the 2 at index 3 becomes the first distinct even number. Balance = .
- Wait: In Example 3,
[2, 3, 2]is balanced because distinct even is{2}and distinct odd is{3}. Both counts are 1. Our logic tracks this by updating ranges!
Code Blocks
C++
#include <vector>
#include <algorithm>
#include <map>
#include <unordered_set>
using namespace std;
struct Node {
int mn, mx, lazy;
Node() : mn(1e9), mx(-1e9), lazy(0) {}
Node(int mn, int mx, int lazy) : mn(mn), mx(mx), lazy(lazy) {}
};
class SegmentTree {
public:
vector<Node> tree;
int size;
SegmentTree(const vector<int>& v) {
size = v.size();
tree.assign(4 * size, Node());
build(0, 0, size - 1, v);
}
void build(int i, int l, int r, const vector<int>& v) {
if (l == r) {
tree[i] = Node(v[l], v[l], 0);
return;
}
int mid = (l + r) / 2;
build(2 * i + 1, l, mid, v);
build(2 * i + 2, mid + 1, r, v);
tree[i].mn = min(tree[2 * i + 1].mn, tree[2 * i + 2].mn);
tree[i].mx = max(tree[2 * i + 1].mx, tree[2 * i + 2].mx);
}
void push(int i, int l, int r) {
if (tree[i].lazy != 0) {
tree[i].mn += tree[i].lazy;
tree[i].mx += tree[i].lazy;
if (l != r) {
tree[2 * i + 1].lazy += tree[i].lazy;
tree[2 * i + 2].lazy += tree[i].lazy;
}
tree[i].lazy = 0;
}
}
void update(int start, int end, int val, int i, int l, int r) {
push(i, l, r);
if (l > end || r < start) return;
if (l >= start && r <= end) {
tree[i].lazy += val;
push(i, l, r);
return;
}
int mid = (l + r) / 2;
update(start, end, val, 2 * i + 1, l, mid);
update(start, end, val, 2 * i + 2, mid + 1, r);
tree[i].mn = min(tree[2 * i + 1].mn, tree[2 * i + 2].mn);
tree[i].mx = max(tree[2 * i + 1].mx, tree[2 * i + 2].mx);
}
int findRightmostZero(int i, int l, int r) {
push(i, l, r);
if (tree[i].mn > 0 || tree[i].mx < 0) return -1;
if (l == r) return l;
int mid = (l + r) / 2;
int res = findRightmostZero(2 * i + 2, mid + 1, r);
if (res == -1) res = findRightmostZero(2 * i + 1, l, mid);
return res;
}
};
class Solution {
public:
int longestBalanced(vector<int>& nums) {
int n = nums.size();
map<int, int> last;
vector<int> nextPos(n, n);
for (int i = n - 1; i >= 0; i--) {
if (last.count(nums[i])) nextPos[i] = last[nums[i]];
last[nums[i]] = i;
}
vector<int> prefix(n, 0);
unordered_set<int> seen;
int current = 0;
for (int i = 0; i < n; i++) {
if (seen.find(nums[i]) == seen.end()) {
current += (nums[i] % 2 == 0 ? -1 : 1);
seen.insert(nums[i]);
}
prefix[i] = current;
}
SegmentTree st(prefix);
int ans = 0;
int rZero = st.findRightmostZero(0, 0, n - 1);
if (rZero != -1) ans = rZero + 1;
for (int i = 1; i < n; i++) {
int val = (nums[i - 1] % 2 == 0 ? 1 : -1);
st.update(i, nextPos[i - 1] - 1, val, 0, 0, n - 1);
int idx = st.findRightmostZero(0, 0, n - 1);
if (idx >= i) ans = max(ans, idx - i + 1);
}
return ans;
}
};
Python
class SegmentTree:
def __init__(self, data):
self.n = len(data)
self.mn = [0] * (4 * self.n)
self.mx = [0] * (4 * self.n)
self.lazy = [0] * (4 * self.n)
self._build(0, 0, self.n - 1, data)
def _build(self, i, l, r, data):
if l == r:
self.mn[i] = self.mx[i] = data[l]
return
mid = (l + r) // 2
self._build(2 * i + 1, l, mid, data)
self._build(2 * i + 2, mid + 1, r, data)
self.mn[i] = min(self.mn[2 * i + 1], self.mn[2 * i + 2])
self.mx[i] = max(self.mx[2 * i + 1], self.mx[2 * i + 2])
def _push(self, i, l, r):
if self.lazy[i] != 0:
self.mn[i] += self.lazy[i]
self.mx[i] += self.lazy[i]
if l != r:
self.lazy[2 * i + 1] += self.lazy[i]
self.lazy[2 * i + 2] += self.lazy[i]
self.lazy[i] = 0
def update(self, start, end, val, i, l, r):
self._push(i, l, r)
if l > end or r < start:
return
if l >= start and r <= end:
self.lazy[i] += val
self._push(i, l, r)
return
mid = (l + r) // 2
self.update(start, end, val, 2 * i + 1, l, mid)
self.update(start, end, val, 2 * i + 2, mid + 1, r)
self.mn[i] = min(self.mn[2 * i + 1], self.mn[2 * i + 2])
self.mx[i] = max(self.mx[2 * i + 1], self.mx[2 * i + 2])
def find_rightmost_zero(self, i, l, r):
self._push(i, l, r)
if self.mn[i] > 0 or self.mx[i] < 0:
return -1
if l == r:
return l
mid = (l + r) // 2
res = self.find_rightmost_zero(2 * i + 2, mid + 1, r)
if res == -1:
res = self.find_rightmost_zero(2 * i + 1, l, mid)
return res
class Solution:
def longestBalanced(self, nums: list[int]) -> int:
n = len(nums)
last = {}
next_pos = [n] * n
for i in range(n - 1, -1, -1):
if nums[i] in last:
next_pos[i] = last[nums[i]]
last[nums[i]] = i
prefix = [0] * n
seen = set()
curr = 0
for i in range(n):
if nums[i] not in seen:
curr += -1 if nums[i] % 2 == 0 else 1
seen.add(nums[i])
prefix[i] = curr
st = SegmentTree(prefix)
ans = 0
rz = st.find_rightmost_zero(0, 0, n - 1)
if rz != -1: ans = rz + 1
for i in range(1, n):
val = 1 if nums[i-1] % 2 == 0 else -1
st.update(i, next_pos[i-1] - 1, val, 0, 0, n - 1)
idx = st.find_rightmost_zero(0, 0, n - 1)
if idx >= i:
ans = max(ans, idx - i + 1)
return ans
JavaScript
class SegmentTree {
constructor(data) {
this.n = data.length;
this.mn = new Int32Array(4 * this.n);
this.mx = new Int32Array(4 * this.n);
this.lazy = new Int32Array(4 * this.n);
this.build(0, 0, this.n - 1, data);
}
build(i, l, r, data) {
if (l === r) {
this.mn[i] = this.mx[i] = data[l];
return;
}
const mid = Math.floor((l + r) / 2);
this.build(2 * i + 1, l, mid, data);
this.build(2 * i + 2, mid + 1, r, data);
this.mn[i] = Math.min(this.mn[2 * i + 1], this.mn[2 * i + 2]);
this.mx[i] = Math.max(this.mx[2 * i + 1], this.mx[2 * i + 2]);
}
push(i, l, r) {
if (this.lazy[i] !== 0) {
this.mn[i] += this.lazy[i];
this.mx[i] += this.lazy[i];
if (l !== r) {
this.lazy[2 * i + 1] += this.lazy[i];
this.lazy[2 * i + 2] += this.lazy[i];
}
this.lazy[i] = 0;
}
}
update(start, end, val, i, l, r) {
this.push(i, l, r);
if (l > end || r < start) return;
if (l >= start && r <= end) {
this.lazy[i] += val;
this.push(i, l, r);
return;
}
const mid = Math.floor((l + r) / 2);
this.update(start, end, val, 2 * i + 1, l, mid);
this.update(start, end, val, 2 * i + 2, mid + 1, r);
this.mn[i] = Math.min(this.mn[2 * i + 1], this.mn[2 * i + 2]);
this.mx[i] = Math.max(this.mx[2 * i + 1], this.mx[2 * i + 2]);
}
findRightmostZero(i, l, r) {
this.push(i, l, r);
if (this.mn[i] > 0 || this.mx[i] < 0) return -1;
if (l === r) return l;
const mid = Math.floor((l + r) / 2);
let res = this.findRightmostZero(2 * i + 2, mid + 1, r);
if (res === -1) res = this.findRightmostZero(2 * i + 1, l, mid);
return res;
}
}
/**
* @param {number[]} nums
* @return {number}
*/
var longestBalanced = function(nums) {
const n = nums.length;
const nextPos = new Array(n).fill(n);
const last = new Map();
for (let i = n - 1; i >= 0; i--) {
if (last.has(nums[i])) nextPos[i] = last.get(nums[i]);
last.set(nums[i], i);
}
const prefix = new Array(n);
const seen = new Set();
let curr = 0;
for (let i = 0; i < n; i++) {
if (!seen.has(nums[i])) {
curr += (nums[i] % 2 === 0 ? -1 : 1);
seen.add(nums[i]);
}
prefix[i] = curr;
}
const st = new SegmentTree(prefix);
let ans = 0;
let rz = st.findRightmostZero(0, 0, n - 1);
if (rz !== -1) ans = rz + 1;
for (let i = 1; i < n; i++) {
const val = (nums[i - 1] % 2 === 0 ? 1 : -1);
st.update(i, nextPos[i - 1] - 1, val, 0, 0, n - 1);
const idx = st.findRightmostZero(0, 0, n - 1);
if (idx >= i) ans = Math.max(ans, idx - i + 1);
}
return ans;
};
Key Takeaways
- Lazy Propagation: This technique allows us to update an entire range in time by deferring calculations until they are actually needed.
- Dynamic Prefix Sums: By shifting the start index and updating the range between the current element and its next occurrence, we effectively reset the "distinct" status for that number.
- Search in Segment Trees: We can use the minimum and maximum values stored in a node to "prune" our search, only diving into branches that could possibly contain a zero.
Final Thoughts
This problem is a fantastic example of how a standard technique (prefix sums for balance) can be elevated to a "Hard" level by adding a single constraint: distinctness. In real-world software engineering, this mirrors scenarios in stream processing or real-time analytics, where you need to track unique user sessions or distinct events within a sliding window. Mastering the Segment Tree gives you a powerful tool for any system requiring high-performance range queries and updates.
Top comments (0)