結果

問題 No.956 Number of Unbalanced
コンテスト
ユーザー qwewe
提出日時 2025-05-14 13:24:48
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
実行時間 -
コード長 2,859 bytes
コンパイル時間 755 ms
コンパイル使用メモリ 88,192 KB
実行使用メモリ 16,072 KB
最終ジャッジ日時 2025-05-14 13:25:51
合計ジャッジ時間 4,654 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 6
other TLE * 1 -- * 23
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

// Fenwick tree (BIT)
std::vector<int> bit_ft; // Renamed to avoid conflict if bit is used elsewhere
int ft_size;
int ft_offset;

void ft_update(int idx, int delta) {
    idx += ft_offset; // Apply offset
    for (; idx < ft_size; idx = idx | (idx + 1)) {
        bit_ft[idx] += delta;
    }
}

// Query sum in [0, idx] inclusive (after offset)
int ft_query_inclusive(int idx) {
    idx += ft_offset; // Apply offset
    if (idx < 0) return 0; // Querying for negative original index (after offset still <0)
    if (idx >= ft_size) idx = ft_size -1; // Cap at max index if original val too large

    int sum = 0;
    for (; idx >= 0; idx = (idx & (idx + 1)) - 1) {
        sum += bit_ft[idx];
    }
    return sum;
}

// Query sum for values strictly less than 'val'
// This means sum up to 'val-1'
int ft_query_less_than(int val) {
    return ft_query_inclusive(val - 1);
}


int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    int n;
    std::cin >> n;

    std::vector<int> a(n);
    std::map<int, bool> distinct_values; // To iterate over distinct values
    for (int i = 0; i < n; ++i) {
        std::cin >> a[i];
        distinct_values[a[i]] = true;
    }

    long long total_unbalanced_subarrays = 0;
    
    // O(D * N log N) approach
    // For Fenwick tree, prefix sums range from -N to N.
    // Shifted range [0, 2N]. Size 2N+1.
    ft_size = 2 * n + 1;
    ft_offset = n; // Shift values by n to make them non-negative

    std::vector<int> pref_b(n + 1);
    std::vector<int> b_arr(n);

    for (auto const& [val_x, present_flag] : distinct_values) {
        // Construct B array
        for(int k=0; k<n; ++k) {
            if (a[k] == val_x) {
                b_arr[k] = 1;
            } else {
                b_arr[k] = -1;
            }
        }

        // Construct prefix sums of B
        pref_b[0] = 0;
        for (int k = 0; k < n; ++k) {
            pref_b[k+1] = pref_b[k] + b_arr[k];
        }

        // Initialize BIT for this X
        // std::fill(bit_ft.begin(), bit_ft.end(), 0); // More efficient than assign if vector is reused
        bit_ft.assign(ft_size, 0); // Safe, resizes and fills
        
        // Add pref_b[0] (corresponds to empty prefix for i=0 start)
        ft_update(pref_b[0], 1);

        for (int k = 1; k <= n; ++k) { // k is pref_b index, corresponds to subarray A[...k-1] ending
            // We want pref_b[k] > pref_b[i_prev_pref_idx]
            // where i_prev_pref_idx is from 0 to k-1
            // query_less_than(pref_b[k]) counts # of prev_pref_sums that are < pref_b[k]
            total_unbalanced_subarrays += ft_query_less_than(pref_b[k]);
            ft_update(pref_b[k], 1);
        }
    }

    std::cout << total_unbalanced_subarrays << std::endl;

    return 0;
}
0