CodeForces 85D: Sum of Medians の多角的なアプローチと実装

本記事では、CodeForces 85D - Sum of Medians という問題に対する4つの異なる解法を解説します。この問題は、動的な集合に対する要素の追加、削除、および特定の位置にある要素の総和を求めるクエリを効率的に処理することを求めています。

問題概要

空の集合 $S$ に対して、 $Q$ 個のクエリが与えられます。各クエリは以下の3種類のいずれかです。

add x: $x \in [1, 10^9]$ を集合 $S$ に追加する（既に存在しないことが保証される）。
del x: $x \in [1, 10^9]$ を集合 $S$ から削除する（既に存在することが保証される）。
sum: 集合 $S$ の要素を昇順にソートして得られる数列 $a$ において、インデックスが 1-indexed で $i \equiv 3 \pmod 5$ を満たす $a_i$ の総和を出力する。

制約は $Q \le 10^5$ です。値の範囲が広いため、適切なデータ構造を選択する必要があります。

解法1：座標圧縮とセグメント木

まず考えるべきは、値の範囲 $[1, 10^9]$ を縮小するための離散化（座標圧縮）です。すべてのクエリで登場する値を事前に収集し、ソートして連続したインデックスに割り当てます。

セグメント木の各ノードは、担当する区間内にある集合 $S$ の要素の個数 cnt と、要素を昇順に並べた際のインデックス $k \pmod 5$ ごとの総和 sum[k] を保持します。

ノードのマージ（プッシュアップ）処理では、左の子の要素数に基づいて、右の子のインデックスのオフセットを計算する必要があります。左の子の要素数を L_cnt とすると、親ノードの sum[k] は以下のように計算できます。

左の子からの寄与: left.sum[k]
右の子からの寄与: 右の子内でのインデックス $j$ が $(k - L_{cnt}) \pmod 5$ となる要素の和。つまり right.sum[(k - L_cnt + 5) % 5]

これにより、add および del クエリは $O(\log Q)$ で処理でき、sum クエリは根ノードの sum[3] を参照するだけで $O(1)$ で答えが得られます。計算量は $O(Q \log Q)$ です。


#include <bits/stdc++.h>
using namespace std;
using int64 = long long;

struct SegmentTree {
    struct Node {
        int count;
        int64 sum[5];
    };
    
    vector<Node> tree;
    int n;
    
    SegmentTree(int size) : n(size) {
        tree.resize(4 * n, {0, {0, 0, 0, 0, 0}});
    }
    
    void update(int pos, int val, int delta, int idx = 1, int l = 0, int r = -1) {
        if (r == -1) r = n - 1;
        if (l == r) {
            tree[idx].count += delta;
            tree[idx].sum[1] = (tree[idx].count > 0) ? val : 0;
            return;
        }
        int mid = l + (r - l) / 2;
        if (pos <= mid) update(pos, val, delta, idx * 2, l, mid);
        else update(pos, val, delta, idx * 2 + 1, mid + 1, r);
        
        pull(idx);
    }
    
    void pull(int idx) {
        tree[idx].count = tree[idx * 2].count + tree[idx * 2 + 1].count;
        int shift = tree[idx * 2].count;
        for (int i = 0; i < 5; ++i) {
            tree[idx].sum[i] = tree[idx * 2].sum[i];
            int target_idx = (i - shift % 5 + 5) % 5;
            tree[idx].sum[i] += tree[idx * 2 + 1].sum[target_idx];
        }
    }
    
    int64 query() {
        return tree[1].sum[3];
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int Q;
    cin >> Q;
    
    vector<pair<string, int64>> queries(Q);
    vector<int64> values;
    
    for (int i = 0; i < Q; ++i) {
        string op;
        int64 x;
        cin >> op;
        if (op != "sum") {
            cin >> x;
            values.push_back(x);
        } else {
            x = -1;
        }
        queries[i] = {op, x};
    }
    
    // Discretization
    sort(values.begin(), values.end());
    values.erase(unique(values.begin(), values.end()), values.end());
    
    auto get_idx = [&](int64 v) {
        return lower_bound(values.begin(), values.end(), v) - values.begin();
    };
    
    SegmentTree seg(values.size());
    
    for (const auto& q : queries) {
        if (q.first == "add") {
            int idx = get_idx(q.second);
            seg.update(idx, q.second, 1);
        } else if (q.first == "del") {
            int idx = get_idx(q.second);
            seg.update(idx, q.second, -1);
        } else {
            cout << seg.query() << "\n";
        }
    }
    
    return 0;
}

解法2：動的セグメント木

解法1の拡張として、動的セグメント木を使用する方法があります。この方法では、事前の離散化を行わず、値の範囲 $[1, 10^9]$ をそのまま扱います。ノードは必要な場合にのみ動的に生成されます。

構造は解法1と同様ですが、各ノードが左右の子へのポインタを持つ点が異なります。これにより、クエリをオンラインで（すべての入力を前もって知らなくても）処理できます。空間計算量は $O(Q \log 10^9)$ 、時間計算量は $O(Q \log 10^9)$ となります。


#include <bits/stdc++.h>
using namespace std;
using int64 = long long;

mt19937 rng(1337);

struct DynamicSegTree {
    struct Node {
        int64 sum[5];
        int cnt;
        Node *l, *r;
        Node() : cnt(0), l(nullptr), r(nullptr) {
            fill(sum, sum + 5, 0);
        }
    };
    
    const int64 L = 1, R = 1000000000;
    Node* root;
    
    DynamicSegTree() : root(new Node()) {}
    
    void update(int64 x, int delta, Node*& node, int64 l, int64 r) {
        if (!node) node = new Node();
        if (l == r) {
            node->cnt += delta;
            node->sum[1] = (node->cnt > 0) ? x : 0;
            return;
        }
        int64 mid = l + (r - l) / 2;
        if (x <= mid) update(x, delta, node->l, l, mid);
        else update(x, delta, node->r, mid + 1, r);
        
        // Pull up
        node->cnt = (node->l ? node->l->cnt : 0) + (node->r ? node->r->cnt : 0);
        int left_cnt = (node->l ? node->l->cnt : 0);
        for (int i = 0; i < 5; ++i) {
            node->sum[i] = (node->l ? node->l->sum[i] : 0);
            int target = (i - left_cnt % 5 + 5) % 5;
            node->sum[i] += (node->r ? node->r->sum[target] : 0);
        }
    }
    
    void add(int64 x) { update(x, 1, root, L, R); }
    void del(int64 x) { update(x, -1, root, L, R); }
    
    int64 query() { return root ? root->sum[3] : 0; }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int Q;
    cin >> Q;
    
    DynamicSegTree seg;
    
    while (Q--) {
        string op;
        int64 x;
        cin >> op;
        if (op != "sum") {
            cin >> x;
            if (op == "add") seg.add(x);
            else seg.del(x);
        } else {
            cout << seg.query() << "\n";
        }
    }
    return 0;
}

解法3：平方分割

次に、平方分割を用いたアプローチです。集合の要素をソートして保持し、いくつかのブロック（バケット）に分割して管理します。

各ブロックは deque を用いて保持し、要素の挿入・削除に伴う順序の維持を容易にします。また、各ブロック内の $i \pmod 5$ ごとの和もキャッシュしておきます。ブロックサイズが一定の閾値を超えた場合、隣接するブロックへ要素を移動する「再平衡」処理を行います。この操作はならし解析で高速になります。ブロックサイズは定数調整が必要ですが、 $O(\sqrt{Q})$ 程度に設定することで、全体の計算量を $O(Q \sqrt{Q})$ に抑えることが可能です。


#include <bits/stdc++.h>
using namespace std;
using int64 = long long;

struct SqrtDecomposition {
    int BLOCK_SIZE;
    vector<deque<int64>> blocks;
    vector<array<int64, 5>> block_sums;
    set<int64> global_set;
    
    SqrtDecomposition(int q) {
        BLOCK_SIZE = max(1, (int)sqrt(q) * 2); // Tuning block size
        blocks.emplace_back();
        block_sums.emplace_back();
        block_sums.back().fill(0);
    }
    
    void rebuild_block(int idx) {
        block_sums[idx].fill(0);
        int rank = 1;
        for (auto& val : blocks[idx]) {
            block_sums[idx][rank % 5] += val;
            rank++;
        }
    }
    
    void rebalance() {
        int64 total_sum = 0;
        int offset = 0;
        for (int i = 0; i < (int)blocks.size(); ++i) {
            if ((int)blocks[i].size() > BLOCK_SIZE * 2) {
                // Move last element to next block
                int64 val = blocks[i].back();
                blocks[i].pop_back();
                
                // Adjust current block sums
                int rank_in_block = blocks[i].size() + 1;
                block_sums[i][rank_in_block % 5] -= val;
                
                if (i + 1 == (int)blocks.size()) {
                    blocks.emplace_back();
                    block_sums.emplace_back();
                    block_sums.back().fill(0);
                }
                blocks[i+1].push_front(val);
                // Adjust next block sums (shift everything by 1)
                for (int k = 0; k < 5; ++k) {
                    int old_idx = k;
                    int new_idx = (k + 1) % 5;
                    block_sums[i+1][new_idx] = block_sums[i+1][old_idx];
                }
                block_sums[i+1][1 % 5] += val;
            }
        }
    }
    
    void add(int64 val) {
        global_set.insert(val);
        auto it = global_set.lower_bound(val);
        if (it != global_set.begin()) it--;
        int64 target = *it;
        
        int target_block = -1;
        for (int i = 0; i < (int)blocks.size(); ++i) {
            if (!blocks[i].empty() && blocks[i].front() <= target && target <= blocks[i].back()) {
                target_block = i;
                break;
            }
        }
        if (target_block == -1) target_block = 0;
        
        // Insert maintaining order
        auto pos = lower_bound(blocks[target_block].begin(), blocks[target_block].end(), val);
        blocks[target_block].insert(pos, val);
        rebuild_block(target_block);
        rebalance();
    }
    
    void del(int64 val) {
        global_set.erase(val);
        int target_block = -1;
        for (int i = 0; i < (int)blocks.size(); ++i) {
            auto it = lower_bound(blocks[i].begin(), blocks[i].end(), val);
            if (it != blocks[i].end() && *it == val) {
                target_block = i;
                blocks[i].erase(it);
                break;
            }
        }
        if (target_block != -1) {
            rebuild_block(target_block);
            rebalance();
        }
    }
    
    int64 get_sum() {
        int64 ans = 0;
        int offset = 0;
        for (int i = 0; i < (int)blocks.size(); ++i) {
            // We want global index % 5 == 3.
            // Global index of block start is offset + 1.
            // Required local index k such that (offset + 1 + k - 1) % 5 == 3
            // => (offset + k) % 5 == 3
            // => k % 5 == (3 - offset % 5 + 5) % 5
            int target_mod = (3 - offset % 5 + 5) % 5;
            ans += block_sums[i][target_mod];
            offset += blocks[i].size();
        }
        return ans;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int Q;
    cin >> Q;
    SqrtDecomposition decomp(Q);
    
    while (Q--) {
        string op;
        int64 x;
        cin >> op;
        if (op != "sum") {
            cin >> x;
            if (op == "add") decomp.add(x);
            else decomp.del(x);
        } else {
            cout << decomp.get_sum() << "\n";
        }
    }
    return 0;
}

解法4：平衡二分探索木 (Treap)

最後に、平衡二分探索木を用いる方法です。ここでは実装の簡易さから、FHQ-Treap（非再帰的Split/Mergeを行うTreap）を採用します。

各ノードは部分木のサイズ sz と、同様に $i \pmod 5$ ごとの和 sum[5] を保持します。挿入および削除は、値順序ではなく「ランク（何番目か）」に基づいて木をSplitし、マージすることで実現します。これにより、常にソート済みのシーケンスを維持できます。


ノードの更新処理では、左部分木のサイズ L_sz を考慮し、現在のノードが全体の  $L_sz + 1$  番目にあることを利用して、右部分木の和のインデックスをずらします。この構造により、すべての操作が期待  $O(\log Q)$  で完了します。


#include <bits/stdc++.h>
using namespace std;
using int64 = long long;

mt19937 rng(2024);

struct Treap {
    struct Node {
        int priority;
        int value;
        int size;
        int64 sum[5];
        Node *l, *r;
        
        Node(int v) : priority(rng()), value(v), size(1), l(nullptr), r(nullptr) {
            fill(sum, sum + 5, 0);
            sum[1] = v;
        }
    };
    
    Node* root;
    
    int get_size(Node* t) { return t ? t->size : 0; }
    
    void update(Node* t) {
        if (!t) return;
        t->size = 1 + get_size(t->l) + get_size(t->r);
        
        int l_sz = get_size(t->l);
        // Current node index is l_sz + 1 (1-based)
        
        for (int i = 0; i < 5; ++i) {
            t->sum[i] = (t->l ? t->l->sum[i] : 0);
            int target_r_idx = (i - l_sz - 1 % 5 + 5) % 5;
            t->sum[i] += (t->r ? t->r->sum[target_r_idx] : 0);
        }
        // Add current node value
        t->sum[(l_sz + 1) % 5] += t->value;
    }
    
    void split(Node* t, int k, Node*& a, Node*& b) {
        // Split first k nodes into a, rest into b
        if (!t) {
            a = b = nullptr;
            return;
        }
        if (get_size(t->l) < k) {
            split(t->r, k - get_size(t->l) - 1, t->r, b);
            a = t;
        } else {
            split(t->l, k, a, t->l);
            b = t;
        }
        update(t);
    }
    
    void merge(Node*& t, Node* a, Node* b) {
        if (!a || !b) {
            t = a ? a : b;
        } else if (a->priority > b->priority) {
            merge(a->r, a->r, b);
            t = a;
        } else {
            merge(b->l, a, b->l);
            t = b;
        }
        update(t);
    }
    
    int get_rank(Node* t, int val) {
        if (!t) return 0;
        if (val <= t->value) return get_rank(t->l, val);
        return 1 + get_size(t->l) + get_rank(t->r, val);
    }
    
    void add(int val) {
        int rank = get_rank(root, val);
        Node *l, *r;
        split(root, rank, l, r);
        merge(root, l, new Node(val));
        merge(root, root, r);
    }
    
    void del(int val) {
        int rank = get_rank(root, val);
        Node *l, *m, *r;
        split(root, rank, l, m);
        split(m, 1, m, r);
        delete m; // m is the node to delete
        merge(root, l, r);
    }
    
    int64 query() {
        return root ? root->sum[3] : 0;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    int Q;
    cin >> Q;
    
    Treap treap;
    
    while (Q--) {
        string op;
        int x;
        cin >> op;
        if (op != "sum") {
            cin >> x;
            if (op == "add") treap.add(x);
            else treap.del(x);
        } else {
            cout << treap.query() << "\n";
        }
    }
    return 0;
}


解法の比較
各手法の特性とパフォーマンスの比較は以下の通りです。




手法
オンライン/オフライン
時間計算量
空間計算量




離散化 + セグメント木
オフライン
 $O(Q \log Q)$ 
 $O(Q)$ 


動的セグメント木
オンライン
 $O(Q \log 10^9)$ 
 $O(Q \log 10^9)$ 


平方分割
オンライン
 $O(Q \sqrt{Q})$ 
 $O(Q)$ 


平衡二分探索木 (Treap)
オンライン
 $O(Q \log Q)$  (期待値)
 $O(Q)$ 




定数倍の観点からは、離散化を用いたセグメント木が最も高速に動作する傾向にありますが、平衡二分探索木は柔軟性が高く、オンライン処理が必要な場合に強力です。平方分割は実装が直感的ですが、制約が厳しい場合には定数倍の最適化が重要になります。

手法	オンライン/オフライン	時間計算量	空間計算量
離散化 + セグメント木	オフライン	$O(Q \log Q)$	$O(Q)$
動的セグメント木	オンライン	$O(Q \log 10^9)$	$O(Q \log 10^9)$
平方分割	オンライン	$O(Q \sqrt{Q})$	$O(Q)$
平衡二分探索木 (Treap)	オンライン	$O(Q \log Q)$ (期待値)	$O(Q)$



    

        タグ:
            SegmentTree
            BalancedBinarySearchTree
            SquareRootDecomposition
            Treap
            DataStructures
        
        6月3日 16:44 投稿

異端開発室