AVL木で問題を解いてみる

AVL木を実装してみるの続きです。

平衡二分探索木であるAVL木を使えば「集合に含まれている値の最小値（最大値）」や「集合に含まれている要素のうち x 以上である最小値（x 以下である最大値）」を求める処理を高速に実行することができるようになります。

Contents

1 D – Cutting Woods
2 E – Best Performances
3 E – Least Elements

D – Cutting Woods

問題の趣旨は以下のとおりです。

1メートルごとに切れ目がついたLメートルの木がある。
クエリ 1 では x メートル部分で木を切断する。
クエリ 2 では x メートル部分を含む木の長さを出力せよ。

最初に木の端の座標である 0 と L をAVL木に格納しておきます。

木が切断されたらAVL木にその部分の座標を格納します。

木の長さを求めるときは、入力された値を x とすると、AVL木に格納されている値のうち x 以上である最小値から x 以下である最大値を引けばよいです。

class Program
{
    static (int, int) ReadInt2()
    {
        string[] vs = Console.ReadLine().Split();
        return (int.Parse(vs[0]), int.Parse(vs[1]));
    }

    static void Main()
    {
        (int L, int Q) = ReadInt2();

        AvlTree avl = new AvlTree();
        avl.Add(0);
        avl.Add(L);

        for (int i = 0; i < Q; i++)
        {
            (int c, int x) = ReadInt2();
            if(c == 1)
                avl.Add(x);
            if (c == 2)
            {
                long left = avl.GetUnders(x, false, 1)[0];
                long right = avl.GetOvers(x, false, 1)[0];
                Console.WriteLine(right - left);
            }
        }
    }
}

class Program

{

static (int, int) ReadInt2()

{

string[] vs = Console.ReadLine().Split();

return (int.Parse(vs[0]), int.Parse(vs[1]));

}

static void Main()

{

(int L, int Q) = ReadInt2();

AvlTree avl = new AvlTree();

avl.Add(0);

avl.Add(L);

for (int i = 0; i < Q; i++)

{

(int c, int x) = ReadInt2();

if(c == 1)

avl.Add(x);

if (c == 2)

{

long left = avl.GetUnders(x, false, 1)[0];

long right = avl.GetOvers(x, false, 1)[0];

Console.WriteLine(right - left);

}

AviTreeクラスの定義は AVL木を実装してみるを参照してください。

E – Best Performances

問題の趣旨は以下のとおりです。

すべての要素が 0 で長さが N の配列がある。
Q 回にわたって要素のひとつが更新されるので、その値が大きいもの K 個の総和を計算し、これを出力せよ。

実際に値が大きいもの K 個がどれなのかを調べて総和を計算していたのでは制限時間以内に間に合わせることはできません。

そこで配列の要素を上位 K 個とそれ以外のもの（high, low）にわけます。そして値が a から b に更新されたら

① a が high にある場合
a を high から削除
b が low_max 以上なら high に追加
そうでないなら low に追加して low_max を high に移動

② a が low にある場合
a を low から削除
b が high_min 以下なら low に追加
そうでないなら high に追加後、high_min を low に移動

③ 最初はすべての要素が 0 なので総和は 0。high に値が追加、削除されるたびに差分から総和を更新してこれを出力する。

class Program
{
    static void Main()
    {
        string[] vs = Console.ReadLine().Split();
        int N = int.Parse(vs[0]);
        int K = int.Parse(vs[1]);
        int Q = int.Parse(vs[2]);

        AvlTree high = new AvlTree();
        AvlTree low = new AvlTree();
        int[] A = new int[N];

        for (int i = 0; i < K; i++)
            high.Add(0);
        for (int i = K; i < N; i++)
            low.Add(0);

        long sum = 0;
        List<long> ans = new List<long>();
        for (int i = 0; i < Q; i++)
        {
            vs = Console.ReadLine().Split();
            int x = int.Parse(vs[0]);
            int y = int.Parse(vs[1]);

            x--;

            int oldValue = A[x];
            int newValue = y;

            A[x] = newValue;

            // oldValueがhighにある場合
            //     oldValueをhighから削除
            //     newValueが low_max以上ならhighに追加
            //     そうでないならlowに追加してlow_maxをhighに移動
            // oldValueがLowにある場合
            //     oldValueをlowから削除
            //     newValueが high_min以下ならlowに追加
            //     そうでないならhighに追加後、high_minをlowに移動

            if (high.Search(oldValue))
            {
                high.Remove(oldValue);
                sum -= oldValue;

                long low_max = low.Max();
                if (newValue >= low_max)
                {
                    high.Add(newValue);
                    sum += newValue;
                }
                else
                {
                    low.Add(newValue);
                    low.Remove(low_max);
                    high.Add(low_max);
                    sum += low_max;
                }
            }
            else
            {
                low.Remove(oldValue);

                long high_min = high.Min();
                if (newValue <= high_min)
                {
                    low.Add(newValue);
                }
                else
                {
                    high.Add(newValue);
                    sum += newValue;
                    high.Remove(high_min);
                    sum -= high_min;
                    low.Add(high_min);
                }
            }
            ans.Add(sum);
        }
        foreach (var v in ans)
            Console.WriteLine(v);
    }
}

class Program

{

static void Main()

{

string[] vs = Console.ReadLine().Split();

int N = int.Parse(vs[0]);

int K = int.Parse(vs[1]);

int Q = int.Parse(vs[2]);

AvlTree high = new AvlTree();

AvlTree low = new AvlTree();

int[] A = new int[N];

for (int i = 0; i < K; i++)

high.Add(0);

for (int i = K; i < N; i++)

low.Add(0);

long sum = 0;

List<long> ans = new List<long>();

for (int i = 0; i < Q; i++)

{

vs = Console.ReadLine().Split();

int x = int.Parse(vs[0]);

int y = int.Parse(vs[1]);

x--;

int oldValue = A[x];

int newValue = y;

A[x] = newValue;

// oldValueがhighにある場合

// oldValueをhighから削除

// newValueが low_max以上ならhighに追加

// そうでないならlowに追加してlow_maxをhighに移動

// oldValueがLowにある場合

// oldValueをlowから削除

// newValueが high_min以下ならlowに追加

// そうでないならhighに追加後、high_minをlowに移動

if (high.Search(oldValue))

{

high.Remove(oldValue);

sum -= oldValue;

long low_max = low.Max();

if (newValue >= low_max)

{

high.Add(newValue);

sum += newValue;

}

else

{

low.Add(newValue);

low.Remove(low_max);

high.Add(low_max);

sum += low_max;

}

else

{

low.Remove(oldValue);

long high_min = high.Min();

if (newValue <= high_min)

{

low.Add(newValue);

}

else

{

high.Add(newValue);

sum += newValue;

high.Remove(high_min);

sum -= high_min;

low.Add(high_min);

}

ans.Add(sum);

}

foreach (var v in ans)

Console.WriteLine(v);

}

E – Least Elements

これも前問と似たような問題です。長さ M の区間を前から後ろにずらしながら区間内の下位 K 個の総和を求めよという問題です。

最初の K 個はなにも考えずにそのまま low に入れます。

M – K 個はlow_maxよりも小さければ low に入れ、代わりに low_max を high に移動させます。そうでなければそのまま high に格納します。

最後の N – M 個は上と同じ要領で追加処理をおこなったあと、A[i – M]の削除処理をおこないます。削除対象が low と high のどちらに格納されているかを調べて削除します。もし low から削除した場合は high_min を high から low へ移動させます。

class Program
{
    static void Main()
    {
        string[] vs = Console.ReadLine().Split();
        int N = int.Parse(vs[0]);
        int M = int.Parse(vs[1]);
        int K = int.Parse(vs[2]);
        int[] A = Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray();

        AvlTree high = new AvlTree();
        AvlTree low = new AvlTree();

        List<long> ans = new List<long>();

        long sum = 0;

        for (int i = 0; i < K; i++)
        {
            low.Add(A[i]);
            sum += A[i];
        }
        for (int i = K; i < M; i++)
        {
            long low_max = low.Max();
            if (A[i] < low_max)
            {
                low.Add(A[i]);
                low.Remove(low_max);
                high.Add(low_max);
                sum += A[i] - low_max;
            }
            else
            {
                high.Add(A[i]);
            }
        }
        ans.Add(sum);

        for (int i = M; i < N; i++)
        {
            long low_max = low.Max();
            if (A[i] < low_max)
            {
                low.Add(A[i]);
                low.Remove(low_max);
                high.Add(low_max);
                sum += A[i] - low_max;
            }
            else
            {
                high.Add(A[i]);
            }

            if (low.Search(A[i - M]))
            {
                low.Remove(A[i - M]);
                long high_min = high.Min();
                low.Add(high_min);
                high.Remove(high_min);
                sum += high_min - A[i - M];
            }
            else
            {
                high.Remove(A[i - M]);
            }
            ans.Add(sum);
        }
        Console.WriteLine(string.Join(" ", ans));
    }
}

class Program

{

static void Main()

{

string[] vs = Console.ReadLine().Split();

int N = int.Parse(vs[0]);

int M = int.Parse(vs[1]);

int K = int.Parse(vs[2]);

int[] A = Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray();

AvlTree high = new AvlTree();

AvlTree low = new AvlTree();

List<long> ans = new List<long>();

long sum = 0;

for (int i = 0; i < K; i++)

{

low.Add(A[i]);

sum += A[i];

}

for (int i = K; i < M; i++)

{

long low_max = low.Max();

if (A[i] < low_max)

{

low.Add(A[i]);

low.Remove(low_max);

high.Add(low_max);

sum += A[i] - low_max;

}

else

{

high.Add(A[i]);

}

ans.Add(sum);

for (int i = M; i < N; i++)

{

long low_max = low.Max();

if (A[i] < low_max)

{

low.Add(A[i]);

low.Remove(low_max);

high.Add(low_max);

sum += A[i] - low_max;

}

else

{

high.Add(A[i]);

}

if (low.Search(A[i - M]))

{

low.Remove(A[i - M]);

long high_min = high.Min();

low.Add(high_min);

high.Remove(high_min);

sum += high_min - A[i - M];

}

else

{

high.Remove(A[i - M]);

}

ans.Add(sum);

}

Console.WriteLine(string.Join(" ", ans));

}

D – Cutting Woods

E – Best Performances

E – Least Elements

鳩でも分かるC#管理人からのお願い

コメントについて

管理人のモチベーションアップのために

コメントをどうぞ