수열과 쿼리 8 (백준 13553)

안녕하세요. 오늘은 K이하를 찾을거예요.

문제

https://www.acmicpc.net/problem/13553

아이디어

수들의 값에는 변화가 없으므로 기본적으로 모스를 생각할 수 있습니다.
그러면 문제가 쉬워집니다.

값의 범위가 1부터 10만까지이므로 압축 없이도 세그트리를 쓸 수 있습니다. 트리에 값에는 배열에서 그 값이 몇개들어있는지를 체크하면 됩니다. 모스를 할때 어떤 범위에 수가 추가가 되면 그 수 +-K의 범위 안에 수가 몇개가 있는지 확인하면 됩니다. 이거는 그냥 세그트리로 가능합니다.

소스코드

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#define ll long long
using namespace std;

ll tree[101010] = { 0 }, N;
void change(int idx, int value)
{
    while (idx <= 100000)
    {
        tree[idx] += value;
        idx += idx & -idx;
    }
}
int SUM(int idx)
{
    int ans = 0;
    while (idx > 0)
    {
        ans += tree[idx];
        idx -= idx & -idx;
    }
    return ans;
}

ll arr[101010] = { 0 }, ans = 0, sqrt_N, K;
void PLUS(ll idx)
{
    ans += SUM(min(arr[idx] + K, (ll)(100000))) - SUM(max(arr[idx] - K - 1, (ll)(0)));
    change(arr[idx], 1);
}
void MINUS(ll idx)
{
    change(arr[idx], -1);
    ans -= SUM(min(arr[idx] + K, (ll)(100000))) - SUM(max(arr[idx] - K - 1, (ll)(0)));
}

bool cmp(pair <pair <ll, ll>, ll> A, pair <pair <ll, ll>, ll> B)
{
    if (A.first.first / sqrt_N != B.first.first / sqrt_N) return A.first.first < B.first.first;
    return A.first.second < B.first.second;
}

ll Ans[101010] = { 0 };
int main(void)
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll M, i, a, b;
    vector <pair <pair <ll, ll>, ll> > v;

    cin >> N >> K; sqrt_N = sqrt(N);
    for (i = 1; i <= N; i++) cin >> arr[i];

    cin >> M;
    for (i = 1; i <= M; i++)
    {
        cin >> a >> b;
        v.push_back({ {a,b},i });
    }
    sort(v.begin(), v.end(), cmp);

    ll s = 1, e = 1; PLUS(1);
    for (i = 0; i < M; i++)
    {
        pair <ll, ll> temp = v[i].first;
        while (s > temp.first) PLUS(--s);
        while (e < temp.second) PLUS(++e);
        while (s < temp.first) MINUS(s++);
        while (e > temp.second) MINUS(e--);

        Ans[v[i].second] = ans;
    }

    for (i = 1; i <= M; i++)
        cout << Ans[i] << "\n";
}

감사합니다.