部分配列内の異なる要素数をクエリで求める方法 | C++でのセグメント木とsetによる実装

この記事では、サイズ n の配列 arr[] が与えられ、各クエリが2つの値 (L, R) を持つという問題を扱います。目標は、部分配列に含まれる異なる(ユニークな)要素の数を答えるプログラムを作成することです。

問題の説明

ここで求めるのは、インデックス L-1 から R-1 までの部分配列内に存在する、重複しない整数の合計数です。

具体例で理解しよう

入力

arr[] = {4, 6, 1, 3, 1, 6, 5}
query = [1, 4]

出力

4

説明

クエリ1: L = 1、R = 4 の場合、インデックス0から3までの範囲({4, 6, 1, 3})に含まれる異なる要素の数は 4 です。

クエリ2: L = 2、R = 6 の場合、インデックス1から5までの範囲({6, 1, 3, 1, 6})に含まれる異なる要素は {6, 1, 3} の3つです。

解決アプローチ

シンプルな方法:線形走査 + set

最も単純な解法は、各クエリごとに配列を L から R まで走査し、要素を set に順次挿入していく方法です。set は重複を自動的に排除するため、最終的な set のサイズがそのまま答えになります。ただし、この方法では1回のクエリに O(R − L) の時間がかかるため、クエリ数や配列サイズが大きくなると非効率になります。

効率的な方法:セグメント木を使う

より効果的なのがセグメント木(segment tree)というデータ構造を利用する方法です。あらかじめ各区間の情報を構築しておくことで、指定された範囲の異なる要素数を高速に取得できます。

セグメント木とは、情報を「区間(セグメント)」の単位で管理する特殊な木構造です。葉ノードは配列の各要素に対応し、非葉ノードはその子ノードが表す区間全体の値(ここでは異なる要素の集合)を保持します。本実装では、区間ごとの要素集合の管理に set を使用します。

C++による実装例

#include <bits/stdc++.h>
using namespace std;
set<int>* segmentTree;

// セグメント木の構築
void CreateSegmentTree(int i, int s, int e, int arr[]) {
    if (s == e) {
        segmentTree[i].insert(arr[s]);
        return;
    }
    CreateSegmentTree(2 * i, s, (s + e) / 2, arr);
    CreateSegmentTree(1 + 2 * i, 1 + (s + e) / 2, e, arr);
    segmentTree[i].insert(segmentTree[2 * i].begin(), segmentTree[2 * i].end());
    segmentTree[i].insert(segmentTree[2 * i + 1].begin(), segmentTree[2 * i + 1].end());
}

// 範囲 [a, b] 内の異なる要素の集合を取得
set<int> findDistSubarray(int node, int l, int r, int a, int b) {
    set<int> left, right, distinctSubarray;
    if (b < l || a > r)
        return distinctSubarray;
    if (a <= l && r <= b)
        return segmentTree[node];
    left = findDistSubarray(2 * node, l, (l + r) / 2, a, b);
    distinctSubarray.insert(left.begin(), left.end());
    right = findDistSubarray(1 + 2 * node, 1 + (l + r) / 2, r, a, b);
    distinctSubarray.insert(right.begin(), right.end());
    return distinctSubarray;
}

int main() {
    int arr[] = {4, 6, 1, 3, 1, 6, 5};
    int n = sizeof(arr) / sizeof(arr[0]);
    int query[] = {1, 4};
    int i = (int)ceil(log2(n));
    i = (2 * (pow(2, i))) - 1;
    segmentTree = new set<int>[i];
    CreateSegmentTree(1, 0, n - 1, arr);
    set<int> distCount = findDistSubarray(1, 0, n - 1, (query[0]-1), (query[1]-1));
    cout<<"The number of distinct elements in the subarray is "<<distCount.size();
    return 0;
}

出力

The number of distinct elements in the subarray is 4

まとめ

このように、セグメント木を用いることで、範囲クエリを木構造上で分割統治的に処理できます。単純な走査方式と比べ、同じ配列に対して多数のクエリを処理する場面で大きなメリットがあります。なお、set を用いた実装では区間マージのコストが発生するため、さらに大規模なデータを扱う場合は Mo's algorithm(平方分割)などの手法も検討するとよいでしょう。