CF 1058202023_2A - An Array and Medians of Subarrays

Rating: -
Tags: -
Solve time: 1m 2s
Verified: yes

Solution

Problem Understanding

We are given an array of even length. We want to cut it into contiguous subarrays, every subarray must have odd length, and all of those subarrays must have exactly the same median.

The median here is the usual median of an odd-length array, the middle element after sorting.

The input consists of a single array. The output is simply whether such a partition exists.

The array length can reach $2 \cdot 10^5$. Any approach that explicitly examines all odd-length subarrays is immediately ruled out. There are $O(n^2)$ subarrays, and even computing a median for one subarray is not constant time. We need something close to $O(n \log n)$.

The first subtle point is that the common median is not arbitrary.

Consider a valid partition whose common median is $x$. Every odd-length segment with median $x$ contains at least $\lfloor len/2 \rfloor + 1$ elements that are $\le x$, and also at least $\lfloor len/2 \rfloor + 1$ elements that are $\ge x$.

When several such odd segments are merged together, these guarantees add up. As a consequence, the median value shared by all segments must also be the median of the entire collection obtained by merging them. This observation is the key structural property of the problem.

A second easy-to-miss detail comes from parity. The total length is even, while every segment length is odd. The number of segments must therefore be even.

For example:

n = 2
a = [1, 2]

The only possible partition is [1] | [2]. Their medians are different, so the answer is No.

Another useful example is:

n = 6
a = [1, 2, 3, 3, 2, 1]

Splitting into [1,2,3] and [3,2,1] gives median 2 in both segments, so the answer is Yes.

A case that defeats naive reasoning is:

n = 6
a = [1, 2, 1, 3, 2, 3]

The value 2 looks promising because it is in the middle of the sorted order, yet no valid partition exists. The answer is No.

Approaches

A brute-force solution would enumerate every odd split position and every possible partition, then compute medians of all resulting segments and check whether they match.

Even if we only tried every partition into two odd segments, we would still need median queries on many ranges. Computing medians by sorting each range would cost $O(n \log n)$ per split, leading to roughly $O(n^2 \log n)$ work.

The breakthrough comes from understanding the structure of any valid partition.

Suppose a valid partition exists and every segment has median $x$.

Because the number of segments is even, take the first segment as one group and merge all remaining segments into a second group. The second group is a merge of an odd number of valid segments, so it also has median $x$. The first group obviously has median $x$.

This means:

A valid partition exists if and only if there is an odd index $k$ such that

median(a[1..k])     = x
median(a[k+1..n])   = x

where $x$ is the median value forced by the whole array structure.

For an even-length array, that value is the lower middle element of the globally sorted array, namely the element at position $n/2$ in 1-based indexing. This is exactly the value used in accepted solutions.

The problem is now reduced to:

For every odd split position $k$, determine the median of the prefix and the median of the suffix.

This becomes a range k-th order statistic problem. With coordinate compression and a persistent segment tree, we can obtain the median of any range in $O(\log n)$.

There are only $O(n)$ candidate split positions, so the full complexity becomes $O(n \log n)$.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	$O(n^2 \log n)$ or worse	$O(n)$	Too slow
Persistent Segment Tree	$O(n \log n)$	$O(n \log n)$	Accepted

Algorithm Walkthrough

Read the array and create a sorted copy.
Let $X$ be the element at position $n/2$ in the sorted array using 1-based indexing. This is the only possible common median value.
Coordinate-compress the array values so that they lie in the range $[1, m]$.
Build a persistent segment tree over prefixes.

The root for position $i$ stores frequencies of values in the prefix $a[1..i]$. 5. Implement a range k-th statistic query.

Using roots corresponding to prefix $r$ and prefix $l-1$, we can obtain the k-th smallest value in $a[l..r]$. 6. For every odd split position $k$:

Compute the median of $a[1..k]$.

Compute the median of $a[k+1..n]$.

If both medians equal $X$, output "Yes". 7. If no split works, output "No".

Why it works

Every valid partition has an even number of odd-length segments because the total length is even.

If all segments have median $X$, then merging any odd number of consecutive valid segments still produces a segment with median $X$. Applying this to all segments except the first gives a partition into exactly two odd-length parts, both with median $X$.

The converse is immediate. If there exists a split into two odd-length parts whose medians are both $X$, that split itself is already a valid partition.

So checking all odd split positions is both necessary and sufficient.

Python Solution

import sys
input = sys.stdin.readline

n = int(input())
a = list(map(int, input().split()))

tmp = sorted(a)
target = tmp[n // 2 - 1]

vals = sorted(set(a))
comp = {v: i + 1 for i, v in enumerate(vals)}
m = len(vals)

left = [0]
right = [0]
cnt = [0]
roots = [0] * (n + 1)

def update(prev, l, r, pos):
    node = len(cnt)
    left.append(left[prev])
    right.append(right[prev])
    cnt.append(cnt[prev] + 1)

    if l != r:
        mid = (l + r) // 2
        if pos <= mid:
            left[node] = update(left[prev], l, mid, pos)
        else:
            right[node] = update(right[prev], mid + 1, r, pos)
    return node

for i in range(1, n + 1):
    roots[i] = update(roots[i - 1], 1, m, comp[a[i - 1]])

def kth(u, v, l, r, k):
    if l == r:
        return l

    mid = (l + r) // 2
    left_count = cnt[left[v]] - cnt[left[u]]

    if k <= left_count:
        return kth(left[u], left[v], l, mid, k)
    return kth(right[u], right[v], mid + 1, r, k - left_count)

def median_value(l, r):
    length = r - l + 1
    k = (length + 1) // 2
    idx = kth(roots[l - 1], roots[r], 1, m, k)
    return vals[idx - 1]

for split in range(1, n, 2):
    if median_value(1, split) == target and median_value(split + 1, n) == target:
        print("Yes")
        sys.exit()

print("No")

The coordinate compression step allows the persistent segment tree to work on a compact value range.

Each version of the tree corresponds to a prefix. To obtain frequencies inside a range, we subtract the counts stored in two versions. This is the standard persistent-segment-tree trick for range order statistics.

The median query asks for the $(len+1)/2$-th smallest element because every range we examine has odd length.

A common source of mistakes is the target median. Since the whole array length is even, we use the lower middle element of the globally sorted array, which corresponds to position $n/2$ in 1-based indexing.

Worked Examples

Example 1

Input:

6
1 2 3 3 2 1

Sorted array:

1 1 2 2 3 3

Target median:

X = 2

Split k	Prefix	Prefix Median	Suffix	Suffix Median
1	[1]	1	[2,3,3,2,1]	2
3	[1,2,3]	2	[3,2,1]	2

At $k=3$, both medians equal $2$, so the answer is Yes.

This example demonstrates the reduction to two segments. Once one valid split is found, no further work is needed.

Example 2

Input:

6
1 2 1 3 2 3

Sorted array:

1 1 2 2 3 3

Target median:

X = 2

Split k	Prefix Median	Suffix Median
1	1	2
3	1	3
5	2	3

No odd split produces median $2$ on both sides.

The answer is No.

This example shows that having the correct global middle value is not enough. The split position must make both odd parts agree with it.

Complexity Analysis

Measure	Complexity	Explanation
Time	$O(n \log n)$	Building persistent versions and querying all odd splits
Space	$O(n \log n)$	Persistent segment tree nodes

With $n \le 2 \cdot 10^5$, $O(n \log n)$ comfortably fits the limits. The persistent segment tree performs only logarithmic work per update and per median query.

Test Cases

# helper: run solution on input string, return output string
import sys
import io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)

    input = sys.stdin.readline

    n = int(input())
    a = list(map(int, input().split()))

    tmp = sorted(a)
    target = tmp[n // 2 - 1]

    vals = sorted(set(a))
    comp = {v: i + 1 for i, v in enumerate(vals)}
    m = len(vals)

    left = [0]
    right = [0]
    cnt = [0]
    roots = [0] * (n + 1)

    def update(prev, l, r, pos):
        node = len(cnt)
        left.append(left[prev])
        right.append(right[prev])
        cnt.append(cnt[prev] + 1)

        if l != r:
            mid = (l + r) // 2
            if pos <= mid:
                left[node] = update(left[prev], l, mid, pos)
            else:
                right[node] = update(right[prev], mid + 1, r, pos)
        return node

    for i in range(1, n + 1):
        roots[i] = update(roots[i - 1], 1, m, comp[a[i - 1]])

    def kth(u, v, l, r, k):
        if l == r:
            return l
        mid = (l + r) // 2
        lc = cnt[left[v]] - cnt[left[u]]
        if k <= lc:
            return kth(left[u], left[v], l, mid, k)
        return kth(right[u], right[v], mid + 1, r, k - lc)

    def med(l, r):
        length = r - l + 1
        k = (length + 1) // 2
        idx = kth(roots[l - 1], roots[r], 1, m, k)
        return vals[idx - 1]

    for split in range(1, n, 2):
        if med(1, split) == target and med(split + 1, n) == target:
            return "Yes\n"

    return "No\n"

# provided samples
assert run("4\n1 1 1 1\n") == "Yes\n"
assert run("6\n1 2 3 3 2 1\n") == "Yes\n"
assert run("6\n1 2 1 3 2 3\n") == "No\n"

# custom cases
assert run("2\n1 2\n") == "No\n"
assert run("2\n5 5\n") == "Yes\n"
assert run("8\n7 7 7 7 7 7 7 7\n") == "Yes\n"
assert run("4\n1 2 2 3\n") == "No\n"

Test input	Expected output	What it validates
`2 / 1 2`	`No`	Smallest impossible case
`2 / 5 5`	`Yes`	Smallest possible case
All values equal	`Yes`	Repeated values and many valid splits
`1 2 2 3`	`No`	Boundary parity and median handling

Edge Cases

Consider:

2
1 2

The only odd split is after the first element.

The prefix median is 1, the suffix median is 2.

No split satisfies the condition, so the algorithm outputs No.

Now consider:

2
5 5

The target median is 5.

The split after position 1 gives medians 5 and 5, so the algorithm outputs Yes.

Finally, consider:

8
7 7 7 7 7 7 7 7

Every odd-length range has median 7.

The first checked split already succeeds, and the algorithm immediately returns Yes.

These cases verify the parity argument, the handling of duplicated values, and the correctness of the median queries.