ABC455-E Unbalanced ABC Substrings

問題へのリンク

C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("6");
            WillReturn.Add("AABBCC");
            //4
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("6");
            WillReturn.Add("ABCABC");
            //0
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("10");
            WillReturn.Add("ACABCAABAB");
            //17
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long[] GetSplitArr(string pStr)
    {
        return (pStr == "" ? new string[0] : pStr.Split(' ')).Select(pX => long.Parse(pX)).ToArray();
    }

    const long Geta = 100000000;

    static void Main()
    {
        List<string> InputList = GetInputList();
        long N = long.Parse(InputList[0]);
        char[] CharArr = InputList[1].ToCharArray();

        // A-BのAVL木
        var InsAVL_A_B = new AVL_Set_MultiSet<long>();
        InsAVL_A_B.IsMultiSet = true;

        // A-CのAVL木
        var InsAVL_A_C = new AVL_Set_MultiSet<long>();
        InsAVL_A_C.IsMultiSet = true;

        // B-CのAVL木
        var InsAVL_B_C = new AVL_Set_MultiSet<long>();
        InsAVL_B_C.IsMultiSet = true;

        // A=B=Cを判定するAVL木
        var InsAVL_A_B_C = new AVL_Set_MultiSet<long>();
        InsAVL_A_B_C.IsMultiSet = true;

        // A-BのRunSum
        long[] RunSumA_B = new long[CharArr.Length];

        // A-CのRunSum
        long[] RunSumA_C = new long[CharArr.Length];

        // B-CのRunSum
        long[] RunSumB_C = new long[CharArr.Length];

        // A=B=Cを判定するRunSum
        long[] RunSumA_B_C = new long[CharArr.Length];

        for (long I = 0; I <= CharArr.GetUpperBound(0); I++) {
            if (CharArr[I] == 'A') {
                RunSumA_B[I] = 1;
                RunSumA_C[I] = 1;
                RunSumA_B_C[I] = Geta + 1;
            }
            if (CharArr[I] == 'B') {
                RunSumA_B[I] = -1;
                RunSumB_C[I] = 1;
                RunSumA_B_C[I] = -Geta;
            }
            if (CharArr[I] == 'C') {
                RunSumA_C[I] = -1;
                RunSumB_C[I] = -1;
                RunSumA_B_C[I] = -1;
            }

            if (I > 0) {
                RunSumA_B[I] += RunSumA_B[I - 1];
                RunSumA_C[I] += RunSumA_C[I - 1];
                RunSumB_C[I] += RunSumB_C[I - 1];
                RunSumA_B_C[I] += RunSumA_B_C[I - 1];
            }
            InsAVL_A_B.Add(RunSumA_B[I]);
            InsAVL_A_C.Add(RunSumA_C[I]);
            InsAVL_B_C.Add(RunSumB_C[I]);
            InsAVL_A_B_C.Add(RunSumA_B_C[I]);
        }

        long OKVal_A_B = 0;
        long OKVal_A_C = 0;
        long OKVal_B_C = 0;
        long OKVal_A_B_C = 0;

        long Answer = N * (N + 1) / 2;
        for (long I = 0; I <= CharArr.GetUpperBound(0); I++) {
            long MinusCnt = 0;
            long Cnt1 = InsAVL_A_B.GetRangeCnt(OKVal_A_B, OKVal_A_B);
            long Cnt2 = InsAVL_A_C.GetRangeCnt(OKVal_A_C, OKVal_A_C);
            long Cnt3 = InsAVL_B_C.GetRangeCnt(OKVal_B_C, OKVal_B_C);
            long Cnt4 = InsAVL_A_B_C.GetRangeCnt(OKVal_A_B_C, OKVal_A_B_C);
            MinusCnt = Cnt1 + Cnt2 + Cnt3 - 2 * Cnt4;

            if (CharArr[I] == 'A') {
                OKVal_A_B++;
                OKVal_A_C++;
                OKVal_A_B_C += (Geta + 1);
            }
            if (CharArr[I] == 'B') {
                OKVal_A_B--;
                OKVal_B_C++;
                OKVal_A_B_C -= Geta;

            }
            if (CharArr[I] == 'C') {
                OKVal_A_C--;
                OKVal_B_C--;
                OKVal_A_B_C--;
            }
            InsAVL_A_B.Remove(RunSumA_B[I]);
            InsAVL_A_C.Remove(RunSumA_C[I]);
            InsAVL_B_C.Remove(RunSumB_C[I]);
            InsAVL_A_B_C.Remove(RunSumA_B_C[I]);

            Answer -= MinusCnt;
        }
        Console.WriteLine(Answer);
    }
}

#region AVL_Set_MultiSet
/// <summary>
/// 要素の追加、削除、検索、取得が可能な集合を表します．
/// </summary>
/// <typeparam name="T">優先度付きキュー内の要素の型を指定します．</typeparam>
/// <remarks>内部的にはAVL木によって実装されています．</remarks>
internal class AVL_Set_MultiSet<T>
{
    Node root;
    readonly IComparer<T> comparer;
    readonly Node nil;

    /// <summary>
    /// 多重集合かどうかを表します．
    /// </summary>
    internal bool IsMultiSet { get; set; }
    internal AVL_Set_MultiSet(IComparer<T> comparer)
    {
        nil = new Node(default(T));
        root = nil;
        this.comparer = comparer;
    }

    internal AVL_Set_MultiSet() : this(Comparer<T>.Default) { }

    /// <summary>
    /// 要素をコレクションに追加します．
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます．</remarks>
    internal bool Add(T v)
    {
        return insert(ref root, v);
    }

    /// <summary>
    /// v が存在するならコレクションから削除します．
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます．</remarks>
    internal bool Remove(T v)
    {
        return remove(ref root, v);
    }

    /// <summary>
    /// 0-indexed で index 番目の要素をコレクションから取得します．
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます．</remarks>
    internal T this[int index] { get { return find(root, index); } }
    internal int Count { get { return root.Count; } }

    internal void RemoveAt(int k)
    {
        if (k < 0 || k >= root.Count) throw new ArgumentOutOfRangeException();
        removeAt(ref root, k);
    }

    /// <summary>
    /// このコレクションに含まれる要素を昇順に並べて返します．
    /// </summary>
    /// <remarks>この操作は計算量 O(N) で実行されます．</remarks>
    internal T[] Items
    {
        get
        {
            T[] ret = new T[root.Count];
            int k = 0;
            walk(root, ret, ref k);
            return ret;
        }
    }

    private void walk(Node t, T[] a, ref int k)
    {
        if (t.Count == 0) return;
        walk(t.lst, a, ref k);
        a[k++] = t.Key;
        walk(t.rst, a, ref k);
    }

    private bool insert(ref Node t, T key)
    {
        if (t.Count == 0) { t = new Node(key); t.lst = t.rst = nil; t.Update(); return true; }
        int cmp = comparer.Compare(t.Key, key);
        bool res;
        if (cmp > 0)
            res = insert(ref t.lst, key);
        else if (cmp == 0) {
            if (IsMultiSet) res = insert(ref t.lst, key);
            else return false;
        }
        else res = insert(ref t.rst, key);
        balance(ref t);
        return res;
    }

    private bool remove(ref Node t, T key)
    {
        if (t.Count == 0) return false;
        int cmp = comparer.Compare(key, t.Key);
        bool ret;
        if (cmp < 0) ret = remove(ref t.lst, key);
        else if (cmp > 0) ret = remove(ref t.rst, key);
        else {
            ret = true;
            var k = t.lst.Count;
            if (k == 0) { t = t.rst; return true; }
            if (t.rst.Count == 0) { t = t.lst; return true; }

            t.Key = find(t.lst, k - 1);
            removeAt(ref t.lst, k - 1);
        }
        balance(ref t);
        return ret;
    }

    private void removeAt(ref Node t, int k)
    {
        int cnt = t.lst.Count;
        if (cnt < k) removeAt(ref t.rst, k - cnt - 1);
        else if (cnt > k) removeAt(ref t.lst, k);
        else {
            if (cnt == 0) { t = t.rst; return; }
            if (t.rst.Count == 0) { t = t.lst; return; }

            t.Key = find(t.lst, k - 1);
            removeAt(ref t.lst, k - 1);
        }
        balance(ref t);
    }

    private void balance(ref Node t)
    {
        int balance = t.lst.Height - t.rst.Height;
        if (balance == -2) {
            if (t.rst.lst.Height - t.rst.rst.Height > 0) { rotR(ref t.rst); }
            rotL(ref t);
        }
        else if (balance == 2) {
            if (t.lst.lst.Height - t.lst.rst.Height < 0) rotL(ref t.lst);
            rotR(ref t);
        }
        else t.Update();
    }

    private T find(Node t, int k)
    {
        if (k < 0 || k > root.Count) throw new ArgumentOutOfRangeException();
        while (true) {
            if (k == t.lst.Count) return t.Key;
            else if (k < t.lst.Count) t = t.lst;
            else { k -= t.lst.Count + 1; t = t.rst; }
        }
    }
    /// <summary>
    /// コレクションに含まれる要素であって、 v 以上の最小の要素の番号を返します。
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます．</remarks>
    internal int LowerBound(T v)
    {
        // 追加機能 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // 追加機能 v 以上な要素が無い場合は-1を返す
        T MaxVal = this[this.Count - 1];
        int cmp = comparer.Compare(MaxVal, v);
        if (cmp == -1) return -1;

        int k = 0;
        Node t = root;
        while (true) {
            if (t.Count == 0) return k;
            if (comparer.Compare(v, t.Key) <= 0) t = t.lst;
            else { k += t.lst.Count + 1; t = t.rst; }
        }
    }
    /// <summary>
    /// コレクションに含まれる要素であって、 v より真に大きい、最小の要素の番号を返します。
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます．</remarks>
    internal int UpperBound(T v)
    {
        // 追加機能 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // 追加機能 v 超えな要素が無い場合は-1を返す
        T MaxVal = this[this.Count - 1];
        int cmp = comparer.Compare(MaxVal, v);
        if (cmp <= 0) return -1;

        int k = 0;
        Node t = root;
        while (true) {
            if (t.Count == 0) return k;
            if (comparer.Compare(t.Key, v) <= 0) { k += t.lst.Count + 1; t = t.rst; }
            else t = t.lst;
        }
    }

    // 追加機能01 V未満で最大の要素の番号を返す
    internal int Lower_Max(T v)
    {
        // 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // v 未満な要素が無い場合は-1を返す
        T MinVal = this[0];
        int cmp = comparer.Compare(MinVal, v);
        if (cmp >= 0) return -1;

        // v 以上の件数を調べる
        int More_Or_Equal_cnt = this.Count;
        int UB = this.Count - 1;
        int LowerB = LowerBound(v);
        if (IsValidInd(LowerB)) {
            More_Or_Equal_cnt -= (UB - LowerB + 1);
        }
        return More_Or_Equal_cnt - 1;
    }

    // 追加機能02 V以下で最大の要素の番号を返す
    internal int LowerOrEqual_Max(T v)
    {
        // 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // v 以下な要素が無い場合は-1を返す
        T MinVal = this[0];
        int cmp = comparer.Compare(MinVal, v);
        if (cmp > 0) return -1;

        // v 超えの件数を調べる
        int More_cnt = this.Count;
        int UB = this.Count - 1;
        int UpperB = UpperBound(v);
        if (IsValidInd(UpperB)) {
            More_cnt -= (UB - UpperB + 1);
        }
        return More_cnt - 1;
    }

    // 追加機能03 LowerBoundなどで返したIndが、有効範囲かを判定
    internal bool IsValidInd(int pInd)
    {
        if (pInd < 0) return false;
        if (this.Count <= pInd) return false;
        return true;
    }

    // 追加機能04 Vを含むかを返す
    internal bool Contains(T v)
    {
        int LowerB = LowerBound(v);
        if (IsValidInd(LowerB) == false) {
            return false;
        }
        int cmp = comparer.Compare(this[LowerB], v);
        return cmp == 0;
    }

    // 追加機能05 Min以上な値の個数を返す
    internal long GetMoreOrEqualCnt(T pMinVal)
    {
        int ResultInd = LowerBound(pMinVal);
        if (IsValidInd(ResultInd) == false) return 0;
        return Count - 1 - ResultInd + 1;
    }

    // 追加機能06 Min超えな値の個数を返す
    internal long GetMoreStrictCnt(T pMinVal)
    {
        int ResultInd = UpperBound(pMinVal);
        if (IsValidInd(ResultInd) == false) return 0;
        return Count - 1 - ResultInd + 1;
    }

    // 追加機能07 Max以下な値の個数を返す
    internal long GetLessOrEqualCnt(T pMaxVal)
    {
        int ResultInd = LowerOrEqual_Max(pMaxVal);

        if (IsValidInd(ResultInd) == false) return 0;
        return ResultInd + 1;
    }

    // 追加機能08 Max未満な値の個数を返す
    internal long GetLessStrictCnt(T pMaxVal)
    {
        int ResultInd = Lower_Max(pMaxVal);
        if (ResultInd == -1) return 0;
        return ResultInd + 1;
    }

    // 追加機能09 Min以上Max以下な値の個数を返す
    internal long GetRangeCnt(T pMinVal, T pMaxVal)
    {
        int ResultInd1 = LowerBound(pMinVal);
        if (IsValidInd(ResultInd1) == false) return 0;

        int ResultInd2 = LowerOrEqual_Max(pMaxVal);
        if (IsValidInd(ResultInd2) == false) return 0;

        if (ResultInd1 > ResultInd2) return 0;
        return ResultInd2 - ResultInd1 + 1;
    }

    private void rotR(ref Node t)
    {
        Node l = t.lst;
        t.lst = l.rst;
        l.rst = t;
        t.Update();
        l.Update();
        t = l;
    }

    private void rotL(ref Node t)
    {
        Node r = t.rst;
        t.rst = r.lst;
        r.lst = t;
        t.Update();
        r.Update();
        t = r;
    }

    class Node
    {
        internal Node(T key)
        {
            Key = key;
        }
        internal int Count { get; private set; }
        internal int Height { get; private set; }
        internal T Key { get; set; }
        internal Node lst, rst;
        internal void Update()
        {
            Count = 1 + lst.Count + rst.Count;
            Height = 1 + Math.Max(lst.Height, rst.Height);
        }
        public override string ToString()
        {
            return string.Format("Count = {0}, Key = {1}", Count, Key);
        }
    }
}
#endregion

解説

Aの個数 != Bの個数 かつ Aの個数 != Cの個数 かつ Bの個数 != Cの個数
ドモルガンで変形すると
(Aの個数 = Bの個数 または Aの個数 = Cの個数 または Bの個数 = Cの個数) の否定
となります。

Aの個数 = Bの個数 を X
Aの個数 = Cの個数 を Y
Bの個数 = Cの個数 を Z
として個数定理(包除原理)で考えます。
X または Y または Z (要素数) = 
X + Y + Z - XY - XZ - YZ + XYZ
です。
XYZ ⇔ XY
XYZ ⇔ XZ
XYZ ⇔ YZ
なので

X + Y + Z - 2 * XYZ
で求めることができます。

X,Y,Zは最初の累積和をAVL木に入れ差分更新で良いです。
XYZは難しいですが、
Aを 1000000000 + 1
Bを-1000000000
Cを-1
として累積和を求めることで対応可能です。

類題

ABC441-E A > B substring