题解：AtCoder AT_awc0027_e Selection of Contiguous Intervals

【题目来源】

AtCoder：E - Selection of Contiguous Intervals

【题目描述】

Takahashi has a sequence of \(N\) integers \(A_1, A_2, \ldots, A_N\).
高桥有一个包含 \(N\) 个整数的序列 \(A_1, A_2, \ldots, A_N\)。

Takahashi wants to select one contiguous interval from this sequence. Specifically, he chooses a pair of integers \((l, r)\) (\(1 \leq l \leq r \leq N\)) and extracts the elements from the \(l\)-th to the \(r\)-th.
高桥想从这个序列中选择一个连续区间。具体来说，他选择一对整数 \((l, r)\)（\(1 \leq l \leq r \leq N\)），并提取第 \(l\) 到第 \(r\) 个元素。

The score of the chosen interval is defined as follows:
所选区间的得分定义如下：

\(f(l, r) = \sum_{i=l}^{r} A_i + (r - l + 1) \times M\)

That is, the score is the sum of the elements in the interval plus the product of the interval length \((r - l + 1)\) and a positive integer \(M\).
也就是说，得分是区间内元素之和加上区间长度 \((r - l + 1)\) 与一个正整数 \(M\) 的乘积。

Find the number of integer pairs \((l, r)\) (\(1 \leq l \leq r \leq N\)) such that the score is at most \(K\).
求满足得分不超过 \(K\) 的整数对 \((l, r)\)（\(1 \leq l \leq r \leq N\)）的数量。

【输入】

\(N\) \(M\) \(K\)
\(A_1\) \(A_2\) \(\cdots\) \(A_N\)

• The first line contains the length of the sequence \(N\), the score coefficient \(M\), and the score upper bound \(K\), separated by spaces.
• The second line contains the elements of the sequence \(A_1, A_2, \ldots, A_N\), separated by spaces.

【输出】

Print in one line the number of integer pairs \((l, r)\) (\(1 \leq l \leq r \leq N\)) such that the score is at most \(K\).

【输入样例】

5 2 10
1 3 2 4 1

【输出样例】

9

【核心思想】

1. 问题分析：给定长度为 \(N\) 的序列 \(A_i\) 和正整数 \(M\)、\(K\)，求满足 \(f(l, r) = \sum_{i=l}^{r} A_i + (r - l + 1) \times M \leq K\) 的区间 \([l, r]\) 数量。这是一个树状数组 + 离散化问题，关键在于将区间和条件转化为前缀和的比较。

2. 算法选择：

   • 前缀和转化：定义 \(sa[i] = \sum_{j=1}^{i} (A_j + M)\)，则 \(f(l, r) = sa[r] - sa[l-1]\)
   • 条件转化：\(f(l, r) \leq K\) 等价于 \(sa[l-1] \geq sa[r] - K\)
   • 树状数组统计：对于每个 \(r\)，统计满足条件的 \(l-1\) 的数量
   • 离散化：由于 \(sa\) 值域较大，需要离散化处理

3. 关键步骤：

   • 计算前缀和：\(sa[i] = sa[i-1] + A_i + M\)（将每个 \(A_i\) 加上 \(M\)）
   • 离散化处理：
     • 收集所有需要离散化的值：\(sa[i]\) 和 \(sa[i] - K\)（\(i = 0\) 到 \(N\)）
     • 排序去重，建立原值到离散化索引的映射
   • 树状数组统计（遍历 \(r = 1\) 到 \(N\)）：
     • 查询目标：target = sa[r] - K
     • 需要统计已插入的 \(sa[l-1]\) 中满足 \(sa[l-1] \geq target\) 的个数
     • 利用树状数组前缀和性质：ge_target = total - query(target_idx - 1)
     • 累加答案：ans += ge_target
     • 将当前 \(sa[r]\) 插入树状数组：add(mp[sa[r]], 1)
   • 输出 ans

4. 时间/空间复杂度：

   • 时间复杂度：\(O(N \log N)\)，离散化 \(O(N \log N)\)，树状数组操作 \(O(N \log N)\)
   • 空间复杂度：\(O(N)\)，前缀和数组、离散化数组、树状数组

5. 树状数组 + 离散化的核心思想：

   • 前缀和转化：将区间和问题转化为两个前缀和的差，\(f(l, r) = sa[r] - sa[l-1]\)
   • 不等式变形：\(sa[r] - sa[l-1] \leq K \Leftrightarrow sa[l-1] \geq sa[r] - K\)
   • 离线统计：遍历右端点 \(r\)，查询满足条件的左端点 \(l-1\) 的数量
   • 离散化必要性：\(sa\) 值域可能很大（\(|A_i| \leq 10^9\)），需要离散化到 \(O(N)\) 范围才能使用树状数组
   • 树状数组优势：支持单点修改和前缀和查询，时间复杂度 \(O(\log N)\)
   • 适用于区间统计、逆序对、动态前缀和问题

【解题思路】

【算法标签】

树状数组

【代码详解】

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 200005;
int n, m, k;
int a[N], sa[N], tr[N * 2], ans;  // tr数组开2倍大小，因为离散化后最多有2*(n+1)个不同的值
map<int, int> mp;  // 用于存储原值到离散化后索引的映射

// 离散化函数：将原始值映射到1开始的连续整数
vector<int> compress(vector<int> a)
{
    auto b = a;  // 复制一份原始数据
    sort(b.begin(), b.end());  // 排序
    b.erase(unique(b.begin(), b.end()), b.end());  // 去重
    vector<int> res(a.size());  // 创建结果数组，指定大小避免越界
    for (int i = 0; i < a.size(); i++)
    {
        // 使用二分查找获取离散化后的索引（从1开始）
        res[i] = lower_bound(b.begin(), b.end(), a[i]) - b.begin() + 1;
    }
    return res;
}

// 计算x的二进制表示中最低位的1所代表的值
int lowbit(int x)
{
    return x & -x;
}

// 树状数组的更新操作：在位置x增加c
void add(int x, int c, int size)
{
    // 向后更新：更新当前位置及其父节点
    for (int i = x; i <= size; i += lowbit(i))
    {
        tr[i] += c;
    }
}

// 树状数组的查询操作：查询前缀和[1, x]
int query(int x)
{
    int res = 0;
    // 向前查询：累加当前位置及其左子树的和
    for (int i = x; i; i -= lowbit(i))
    {
        res += tr[i];
    }
    return res;
}

signed main()
{
    cin >> n >> m >> k;

    // 读取数组并计算前缀和：sa[i] = Σ(A[j] + m) for j=1..i
    for (int i = 1; i <= n; i++)
    {
        cin >> a[i];
        sa[i] = sa[i - 1] + a[i] + m;  // 计算前缀和，每个元素都加上m
    }

    // 收集所有需要离散化的值
    vector<int> alls;
    alls.push_back(0);  // 添加sa[0] = 0

    // 对于每个前缀和，添加两个值：
    // 1. 前缀和本身sa[i]：用于插入树状数组
    // 2. 前缀和减去k：sa[i]-k：用于查询条件
    for (int i = 0; i <= n; i++)
    {
        alls.push_back(sa[i]);
        alls.push_back(sa[i] - k);
    }

    // 离散化处理
    vector<int> compressed = compress(alls);

    // 建立映射：原值 -> 离散化后的索引
    for (int i = 0; i < alls.size(); i++)
    {
        mp[alls[i]] = compressed[i];
    }

    int sz = compressed.size();  // 离散化后的总大小

    // 初始化树状数组，先添加sa[0] = 0
    add(mp[0], 1, sz);

    // 遍历每个右端点r（对应区间[?, r]）
    for (int r = 1; r <= n; r++)
    {
        // 查询条件：我们需要找到满足 sa[r] - sa[l-1] <= k 的l
        // 等价于：sa[l-1] >= sa[r] - k
        int target = sa[r] - k;  // 查询目标值
        int target_idx = mp[target];  // 目标值的离散化索引

        // 计算已添加的元素个数：sa[0]到sa[r-1]，共r个
        int total = r;

        // 查询小于target的元素个数
        int less_than = query(target_idx - 1);

        // 大于等于target的元素个数 = 总数 - 小于target的个数
        int ge_target = total - less_than;

        // 累加到答案中
        ans += ge_target;

        // 将当前sa[r]添加到树状数组中，供后续查询使用
        add(mp[sa[r]], 1, sz);
    }

    cout << ans << endl;
    return 0;
}

【运行结果】

5 2 10
1 3 2 4 1
9