Median on Segments (Permutations Edition)

Created by LXC on Fri Oct 13 09:49:54 2023

https://codeforces.com/problemset/problem/1005/E1

ranting: 1800

tag: sortings

problem

给出一个长度为n的排列a，包含1到n各1次。

现在求满足$a_l, a_{l+1}, \cdots, a_r$的中位数是m的数对(l,r)的数目。

如果r-l+1的个数为偶数则中位数是升序排序后中间靠左的一个。

solution

先找出m所在位置p。我们可以确定的是$l\le p \le r$。

对于区间$[l,p]$的每个后缀，统计大于m的个数-小于m的个数，并加入集合L中。

对于区间$[p,r]$的每个前缀，统计大于m的个数-小于m的个数，并加入集合R中。

这样问题转化为在集合L和R中各挑出一个元素使得其和为0或1。

这是一个很常规的套路。

将一个集合中各元素$x$用哈希表$mp$将元素值$x$映射到该元素的频次$cnt_x$，即$mp[x] = cnt_x$，然后遍历另一个集合中的各元素$y$，我们要求的是满足$y+x=0或1$的$x$的数目，$mp[0-x]+mp[1-x]$即为当前$y$的贡献。

code

#include <bits/stdc++.h>
#define SINGLE_INPUT
#define ll long long
#define ull unsigned long long
#define N 500005
#define MOD 998244353
using namespace std;

void sol() {
    int n, m, p;
    cin >> n >> m;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
        if (a[i] == m)
            p = i;
    }
    vector<int> d(n);
    for (int i = p - 1; i >= 0; i--) {
        d[i] += d[i + 1];
        if (a[i] > m)
            d[i]++;
        else if (a[i] < m)
            d[i]--;
    }
    for (int i = p + 1; i < n; i++) {
        d[i] += d[i - 1];
        if (a[i] > m)
            d[i]++;
        else if (a[i] < m)
            d[i]--;
    }
    // for (auto i : d) {
    //     cout << i << " ";
    // }
    map<int, int> mp;
    for (int i = 0; i <= p; i++)
        mp[d[i]]++;
    ll ans = 0;
    for (int i = p; i < n; i++) {
        // d[i] + x = 0 or 1
        ans += mp[-d[i]] + mp[1 - d[i]];
    }
    cout << ans << "\n";
}

int main() {
    cout << setprecision(15) << fixed;
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
#ifndef SINGLE_INPUT
    int t;
    cin >> t;
    while (t--) {
        sol();
    }
#else
    sol();
#endif
    return 0;
}