Find the number of sub arrays in the permutation of first N natural numbers such that their median is M in Python

Finding the number of sub-arrays in a permutation where the median equals a specific value M is a classic problem that can be solved efficiently using prefix sums and hash maps.

The key insight is that for a sub-array to have median M, the number of elements less than M should be at most equal to the number of elements greater than M. We track the balance using a running sum where we subtract 1 for elements less than M and add 1 for elements greater than M.

Algorithm Approach

We maintain a balance counter and use a hash map to store frequency of balance values. When we encounter M, we start counting valid sub-arrays by checking previous balance states that would create valid medians.

Step-by-Step Solution

Here's how the algorithm works ?

def count_subarrays_with_median(arr, m):
    n = len(arr)
    balance_map = {0: 1}  # Initialize with balance 0
    has_m = False
    balance = 0
    result = 0
    
    for i in range(n):
        # Update balance based on current element
        if arr[i] < m:
            balance -= 1
        elif arr[i] > m:
            balance += 1
        
        # Check if we've encountered M
        if arr[i] == m:
            has_m = True
        
        if has_m:
            # Count valid subarrays ending at current position
            # For median M, we need balance = 0 or balance = 1
            if balance in balance_map:
                result += balance_map[balance]
            if (balance - 1) in balance_map:
                result += balance_map[balance - 1]
        else:
            # Store current balance for future use
            balance_map[balance] = balance_map.get(balance, 0) + 1
    
    return result

# Test with example
arr = [3, 5, 6, 4, 2]
m = 5
result = count_subarrays_with_median(arr, m)
print(f"Number of subarrays with median {m}: {result}")

Number of subarrays with median 5: 4

How It Works

Let's trace through the example step by step ?

def trace_algorithm(arr, m):
    print(f"Array: {arr}, Target median: {m}")
    print("Step-by-step trace:")
    
    balance_map = {0: 1}
    has_m = False
    balance = 0
    result = 0
    
    for i, val in enumerate(arr):
        print(f"\nStep {i+1}: Processing {val}")
        
        # Update balance
        if val < m:
            balance -= 1
            print(f"  {val} < {m}, balance -= 1 ? {balance}")
        elif val > m:
            balance += 1
            print(f"  {val} > {m}, balance += 1 ? {balance}")
        else:
            print(f"  Found target {m}!")
            has_m = True
        
        if has_m:
            # Count valid subarrays
            count = 0
            if balance in balance_map:
                count += balance_map[balance]
                print(f"  Found balance {balance} in map, adding {balance_map[balance]}")
            if (balance - 1) in balance_map:
                count += balance_map[balance - 1]
                print(f"  Found balance {balance-1} in map, adding {balance_map[balance-1]}")
            result += count
            print(f"  Added {count} subarrays, total: {result}")
        else:
            balance_map[balance] = balance_map.get(balance, 0) + 1
            print(f"  Updated balance_map: {balance_map}")
    
    return result

# Trace the example
trace_algorithm([3, 5, 6, 4, 2], 5)

Array: [3, 5, 6, 4, 2], Target median: 5
Step-by-step trace:

Step 1: Processing 3
  3 < 5, balance -= 1 ? -1
  Updated balance_map: {0: 1, -1: 1}

Step 2: Processing 5
  Found target 5!
  Found balance -1 in map, adding 1
  Found balance -2 in map, adding 1
  Added 1 subarrays, total: 1

Step 3: Processing 6
  6 > 5, balance += 1 ? 0
  Found balance 0 in map, adding 1
  Found balance -1 in map, adding 1
  Added 2 subarrays, total: 3

Step 4: Processing 4
  4 < 5, balance -= 1 ? -1
  Found balance -1 in map, adding 1
  Added 1 subarrays, total: 4

Step 5: Processing 2
  2 < 5, balance -= 1 ? -2
  Added 0 subarrays, total: 4

Valid Subarrays Identified

The four valid subarrays with median 5 are ?

def find_valid_subarrays(arr, m):
    valid_subarrays = []
    n = len(arr)
    
    # Check all possible subarrays
    for i in range(n):
        for j in range(i, n):
            subarray = arr[i:j+1]
            if m in subarray:
                sorted_sub = sorted(subarray)
                median_idx = (len(sorted_sub) - 1) // 2
                if sorted_sub[median_idx] == m:
                    valid_subarrays.append(subarray)
    
    return valid_subarrays

# Find all valid subarrays
arr = [3, 5, 6, 4, 2]
m = 5
valid = find_valid_subarrays(arr, m)

print("Valid subarrays with median 5:")
for i, subarray in enumerate(valid, 1):
    sorted_sub = sorted(subarray)
    print(f"{i}. {subarray} ? sorted: {sorted_sub}, median: {sorted_sub[(len(sorted_sub)-1)//2]}")

Valid subarrays with median 5:
1. [5] ? sorted: [5], median: 5
2. [3, 5, 6] ? sorted: [3, 5, 6], median: 5
3. [5, 6] ? sorted: [5, 6], median: 5
4. [5, 6, 4] ? sorted: [4, 5, 6], median: 5

Time and Space Complexity

Conclusion

This efficient algorithm uses prefix sums and hash maps to count subarrays with a specific median in O(n) time. The key insight is tracking the balance between elements smaller and larger than the target median.