Count Subarrays With Median K

Count Subarrays With Median K - Problem

You are given an array nums of size n consisting of distinct integers from 1 to n and a positive integer k.

Return the number of non-empty subarrays in nums that have a median equal to k.

Note:

The median of an array is the middle element after sorting the array in ascending order. If the array is of even length, the median is the left middle element.
For example, the median of [2,3,1,4] is 2, and the median of [8,4,3,5,1] is 4.
A subarray is a contiguous part of an array.

Input & Output

Example 1 — Basic Case

$ Input: nums = [3,2,1], k = 2

› Output: 3

💡 Note: Three subarrays have median 2: [2] (median=2), [3,2] (sorted=[2,3], median=2), [2,1] (sorted=[1,2], median=1... wait, let me recalculate. [2] median=2, [3,2] sorted=[2,3] median=2, [2,1] sorted=[1,2] median=1. Actually [3,2,1] sorted=[1,2,3] median=2. So subarrays: [2], [3,2], [3,2,1] = 3 total.

Example 2 — Single Element

$ Input: nums = [2,1,3], k = 1

› Output: 3

💡 Note: Three subarrays have median 1: [1] (median=1), [2,1] (sorted=[1,2], median=1), and [1,3] (sorted=[1,3], median=1).

Example 3 — Multiple Valid

$ Input: nums = [1,3,2,4], k = 3

› Output: 2

💡 Note: Two subarrays have median 3: [3] (median=3) and [1,3,2] (sorted=[1,2,3], left-middle median=2... wait, that's wrong. Let me recalculate: [3] median=3, [1,3,2] sorted=[1,2,3] median=2 (left-middle), [3,2] sorted=[2,3] median=2, [1,3,2,4] sorted=[1,2,3,4] median=2. Actually only [3] and [3,2,4] work... let me be more careful.

Constraints

n == nums.length
1 ≤ n ≤ 10⁵
1 ≤ nums[i], k ≤ n
All integers in nums are distinct

Visualization

Tap to expand

Asked in

G Google 15 F Facebook 12 a Amazon 8

The key insight is to transform the median problem into a balance counting problem. Convert elements greater than k to +1, less than k to -1, and k to 0, then count subarrays containing k with sum=0. Best approach uses prefix sums with hash map. Time: O(n), Space: O(n)

Common Approaches

✓ Balance Transformation with Hash Map

⏱️ Time: O(n) Space: O(n)

Find position of k, then transform values: elements > k become +1, elements < k become -1, k becomes 0. Count subarrays containing k with equal +1s and -1s using prefix sum and hash map.

Brute Force - Check All Subarrays

⏱️ Time: O(n³ log n) Space: O(n)

For every possible subarray, extract it, sort it, and check if the median (middle element or left-middle for even length) equals k. This is straightforward but inefficient.

Balance Transformation with Hash Map — Algorithm Steps

Find index of k in the array
Transform: >k → +1, <k → -1, =k → 0
Use prefix sums and hash map to count balanced subarrays containing k

Visualization

Tap to expand

Step-by-Step Walkthrough

Find K

Locate position of target value k

Transform

Convert >k to +1, <k to -1, =k to 0

Count

Count subarrays containing k with sum=0

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

int compare(const void *a, const void *b) {
    return (*(int*)a - *(int*)b);
}

int solution(int* nums, int numsSize, int k) {
    int count = 0;
    
    // Check all possible subarrays
    for (int i = 0; i < numsSize; i++) {
        for (int j = i; j < numsSize; j++) {
            int length = j - i + 1;
            int* subarray = (int*)malloc(length * sizeof(int));
            
            // Copy subarray
            for (int idx = 0; idx < length; idx++) {
                subarray[idx] = nums[i + idx];
            }
            
            // Sort subarray
            qsort(subarray, length, sizeof(int), compare);
            
            // Check median
            int medianIdx = (length - 1) / 2;
            if (subarray[medianIdx] == k) {
                count++;
            }
            
            free(subarray);
        }
    }
    
    return count;
}

void parseArray(const char* str, int* arr, int* size) {
    *size = 0;
    const char* p = str;
    while (*p && *p != '[') p++;
    if (*p == '[') p++;
    while (*p && *p != ']') {
        while (*p == ' ' || *p == ',') p++;
        if (*p == ']' || *p == '\0') break;
        arr[(*size)++] = (int)strtol(p, (char**)&p, 10);
    }
}

int main() {
    char line[10000];
    fgets(line, sizeof(line), stdin);
    
    // Parse array
    int nums[1000];
    int numsSize = 0;
    char* ptr = strchr(line, '[');
    ptr++; // Skip [
    char* token = strtok(ptr, ",]");
    while (token != NULL) {
        nums[numsSize++] = atoi(token);
        token = strtok(NULL, ",]");
    }
    
    int k;
    scanf("%d", &k);
    
    int result = solution(nums, numsSize, k);
    printf("%d\n", result);
    
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Single pass to find k position, then two passes for left and right prefix sums

⚠ Quadratic Growth

Space Complexity

O(n)

Hash map stores prefix sum frequencies

⚡ Linearithmic Space

23.4K Views

Medium Frequency

~30 min Avg. Time

867 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Balance Transformation with Hash Map — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler