結果

問題 No.1079 まお
ユーザー qwewe
提出日時 2025-05-14 12:55:48
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 232 ms / 2,000 ms
コード長 9,333 bytes
コンパイル時間 1,108 ms
コンパイル使用メモリ 103,880 KB
実行使用メモリ 20,864 KB
最終ジャッジ日時 2025-05-14 12:57:22
合計ジャッジ時間 4,740 ms
ジャッジサーバーID
(参考情報)
judge5 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 2
other AC * 30
権限があれば一括ダウンロードができます

ソースコード

diff #

#include <iostream>
#include <vector>
#include <stack>
#include <map>
#include <vector>
#include <numeric>
#include <algorithm>

using namespace std;

// Structure to store indices and prefix sums for a specific value in the array A.
// This helps efficiently query the count and sum of indices within a given range [a, b].
struct IndexData {
    vector<int> indices; // List of 1-based indices where the value appears
    vector<long long> prefix_sums; // prefix_sums[k] = sum of indices[0]...indices[k]

    // Build prefix sums after the indices vector is fully populated.
    void build_prefix_sums() {
        prefix_sums.resize(indices.size());
        if (!indices.empty()) {
            // Calculate prefix sums of the indices. Useful for range sum queries.
            prefix_sums[0] = indices[0]; 
            for (size_t i = 1; i < indices.size(); ++i) {
                prefix_sums[i] = prefix_sums[i - 1] + indices[i];
            }
        }
    }

    // Query the count of indices and the sum of indices for this value within the range [a, b] (inclusive).
    // Returns a pair {count, sum_of_indices}.
    pair<long long, long long> query(int a, int b) {
        // If there are no indices for this value, or the query range is invalid, return {0, 0}.
        if (indices.empty() || a > b) { 
            return {0, 0};
        }

        // Use binary search (lower_bound and upper_bound) to find the range of indices within [a, b].
        // lower_bound finds the first index >= a.
        auto it_l = lower_bound(indices.begin(), indices.end(), a);
        // upper_bound finds the first index > b.
        auto it_r = upper_bound(indices.begin(), indices.end(), b);

        // If it_l == it_r, it means no indices fall within the range [a, b].
        if (it_l == it_r) { 
            return {0, 0};
        }

        // Calculate the 0-based indices in the `indices` vector corresponding to the range found.
        int start_idx = distance(indices.begin(), it_l);
        // it_r points one position past the last element <= b. So the index of the last element is it_r - 1.
        int end_idx = distance(indices.begin(), it_r) - 1; 

        // This check is technically redundant if it_l != it_r, but added for robustness.
        if (start_idx > end_idx) { 
             return {0, 0};
        }

        // Calculate the count of indices in the range.
        long long count = end_idx - start_idx + 1;
        // Calculate the sum of indices using the precomputed prefix sums.
        long long sum_indices = prefix_sums[end_idx];
        if (start_idx > 0) {
            sum_indices -= prefix_sums[start_idx - 1];
        }
        
        return {count, sum_indices};
    }
};


int main() {
    // Use faster I/O operations.
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N; // Number of elements in the array A
    long long K; // Target sum for the endpoints of subarray B
    cin >> N >> K;

    vector<long long> A(N + 1); // Use 1-based indexing for array A. Values can be up to 10^9.
    map<long long, IndexData> val_indices; // Map from value to its IndexData structure.
    
    // Read input array A and populate the map `val_indices`.
    for (int i = 1; i <= N; ++i) {
        cin >> A[i];
        val_indices[A[i]].indices.push_back(i); // Store the 1-based index i.
    }

    // Build prefix sums for each value's index list. This is done after all indices are collected.
    for (auto it = val_indices.begin(); it != val_indices.end(); ++it) {
        it->second.build_prefix_sums();
    }

    // Compute Previous Less or Equal Element (PLEE) and Next Less or Equal Element (NLEE) indices.
    // L[p] = largest index k < p such that A[k] <= A[p]. If none exists, L[p] = 0.
    // R[p] = smallest index k > p such that A[k] <= A[p]. If none exists, R[p] = N + 1.
    // These define the maximal range (L[p], R[p]) where A[p] can be the unique minimum.
    vector<int> L(N + 1); 
    vector<int> R(N + 1); 

    // Compute L (PLEE) using a stack.
    stack<int> st_plee;
    for(int i=1; i<=N; ++i) {
        // Pop elements from stack while the top element's value is strictly greater than A[i].
        while(!st_plee.empty() && A[st_plee.top()] > A[i]) {
            st_plee.pop();
        }
        // The PLEE is the index at the top of the stack, or 0 if the stack is empty.
        L[i] = st_plee.empty() ? 0 : st_plee.top();
        // Handle elements with equal value: if A[top] == A[i], pop the top element.
        // This ensures that for subsequent elements, `i` is considered as the PLEE candidate instead of `st_plee.top()`,
        // effectively keeping the largest index among equal values relevant for the PLEE definition.
        if (!st_plee.empty() && A[st_plee.top()] == A[i]) {
           st_plee.pop(); 
        }
        st_plee.push(i); // Push the current index onto the stack.
    }

    // Compute R (NLEE) using a stack, iterating backwards from N to 1.
    stack<int> st_nlee;
    for (int i = N; i >= 1; --i) {
        // Pop elements from stack while the top element's value is strictly greater than A[i].
        while (!st_nlee.empty() && A[st_nlee.top()] > A[i]) {
            st_nlee.pop();
        }
        // The NLEE is the index at the top of the stack, or N+1 if the stack is empty.
        R[i] = st_nlee.empty() ? N + 1 : st_nlee.top();
        // Handle equal values similarly to PLEE, keeping the smallest index among equal values.
        if (!st_nlee.empty() && A[st_nlee.top()] == A[i]) {
           st_nlee.pop(); 
        }
        st_nlee.push(i); // Push the current index onto the stack.
    }

    long long total_sum_len = 0; // Use long long to store the total sum of lengths, which can be large.

    // Iterate through each element A[p] considering it as the potential unique minimum of a subarray B.
    for (int p = 1; p <= N; ++p) {
        // Define the valid range [i_start, i_end] for the left endpoint `i` and [j_start, j_end] for the right endpoint `j`.
        // For A[p] to be the unique minimum of A[i...j], we must have L[p] < i <= p <= j < R[p].
        int i_start = L[p] + 1;
        int i_end = p;
        int j_start = p;
        int j_end = R[p] - 1;

        // If either the range for i or the range for j is invalid (e.g., start > end), skip this p.
        if (i_start > i_end || j_start > j_end) continue; 

        // Calculate the lengths of the left range (for i) and right range (for j).
        int len_L = i_end - i_start + 1;
        int len_R = j_end - j_start + 1;

        // Apply the "smaller-to-larger" optimization heuristic:
        // Iterate through the smaller range and perform queries on the larger range.
        // This optimization helps improve the overall time complexity.
        if (len_L <= len_R) {
            // Iterate through i in the left range [L[p]+1, p].
            for (int i = i_start; i <= i_end; ++i) {
                long long target_val = K - A[i]; // Calculate the required value A[j] such that A[i] + A[j] = K.
                
                // Check if this target value exists in the array using map::find for efficiency.
                auto it = val_indices.find(target_val);
                if (it != val_indices.end()) {
                    // Query for indices j in the right range [p, R[p]-1] that have the target value.
                    pair<long long, long long> query_res = it->second.query(j_start, j_end);
                    long long count = query_res.first;   // Number of such j's found.
                    long long sum_j = query_res.second; // Sum of these j indices.
                    if (count > 0) {
                        // Add the total length contributed by pairs starting at this `i`.
                        // The length of subarray A[i...j] is j - i + 1.
                        // Summing over all found j's: Sum(j - i + 1) = Sum(j) - Sum(i - 1) = Sum(j) - (i - 1) * count.
                        total_sum_len += sum_j - (long long)(i - 1) * count;
                    }
                }
            }
        } else { // len_L > len_R
            // Iterate through j in the right range [p, R[p]-1].
            for (int j = j_start; j <= j_end; ++j) {
                 long long target_val = K - A[j]; // Calculate the required value A[i] such that A[i] + A[j] = K.
                 
                 // Check if this target value exists.
                 auto it = val_indices.find(target_val);
                 if (it != val_indices.end()) {
                     // Query for indices i in the left range [L[p]+1, p] that have the target value.
                     pair<long long, long long> query_res = it->second.query(i_start, i_end);
                     long long count = query_res.first;  // Number of such i's found.
                     long long sum_i = query_res.second; // Sum of these i indices.
                     if (count > 0) {
                        // Add the total length contributed by pairs ending at this `j`.
                        // Sum(j - i + 1) = Sum(j + 1) - Sum(i) = (j + 1) * count - Sum(i).
                         total_sum_len += (long long)(j + 1) * count - sum_i;
                     }
                 }
            }
        }
    }

    // Output the final computed total sum of lengths.
    cout << total_sum_len << endl;

    return 0;
}
0