ABC469-C Cantrip

問題へのリンク

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("5");
            WillReturn.Add("oxoxo");
            //2
            //4
            //5
            //5
            //5
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("3");
            WillReturn.Add("ooo");
            //3
            //3
            //3
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("1");
            WillReturn.Add("x");
            //1
        }
        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();
    }

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

        var Ins_Fenwick_Tree_Atari = new Fenwick_Tree(UB);
        var Ins_Fenwick_Tree_Hazure = new Fenwick_Tree(UB);

        for (long I = 0; I <= UB; I++) {
            if (CharArr[I] == 'o') Ins_Fenwick_Tree_Atari[I] = 1;
            if (CharArr[I] == 'x') Ins_Fenwick_Tree_Hazure[I] = 1;
        }

        long[] RunSum = new long[UB + 1];
        for (long I = 0; I <= UB; I++) {
            if (CharArr[I] == 'x') RunSum[I] = 1;
            if (I > 0) {
                RunSum[I] += RunSum[I - 1];
            }
        }
        var IndListDict = new Dictionary<long, List<long>>();
        for (long I = 0; I <= UB; I++) {
            if (IndListDict.ContainsKey(RunSum[I]) == false) {
                IndListDict[RunSum[I]] = new List<long>();
            }
            IndListDict[RunSum[I]].Add(I);
        }

        for (long I = 0; I <= UB; I++) {
            if (I == UB) {
                Console.WriteLine(UB + 1);
                continue;
            }

            long AtariCnt = Ins_Fenwick_Tree_Atari.GetSum(0, I);
            long HazureCnt = Ins_Fenwick_Tree_Hazure.GetSum(0, I);
            if (AtariCnt == 0) {
                Console.WriteLine(I + 1);
                continue;
            }

            long StopRunSum = AtariCnt + HazureCnt;
            if (IndListDict.ContainsKey(StopRunSum) == false) {
                Console.WriteLine(UB + 1);
                continue;
            }

            long ResultInd = ExecNibunhou_UpperBound(I, IndListDict[StopRunSum]);
            if (ResultInd == -1) {
                Console.WriteLine(UB + 1);
            }
            else {
                Console.WriteLine(IndListDict[StopRunSum][(int)ResultInd] + 1);
            }
        }
    }

    // 二分法で、Val超えで最小の値を持つ、添字を返す
    static long ExecNibunhou_UpperBound(long pVal, List<long> pList)
    {
        // 要素が0件のケース
        if (pList.Count == 0) return -1;

        // 最後の要素がVal以下の特殊ケース
        if (pVal >= pList.Last()) {
            return -1;
        }
        // 最初の要素がVal超えの特殊ケース
        if (pVal < pList[0]) {
            return 0;
        }

        long L = 0;
        long R = pList.Count - 1;

        while (L + 1 < R) {
            long Mid = (L + R) / 2;

            if (pList[(int)Mid] > pVal) {
                R = Mid;
            }
            else {
                L = Mid;
            }
        }
        return R;
    }
}

// フェニック木
#region Fenwick_Tree
internal class Fenwick_Tree
{
    private long[] mBitArr;
    private long mExternalArrUB;

    // ノードのIndexの列挙を返す
    internal IEnumerable<long> GetNodeIndEnum()
    {
        for (long I = 0; I <= mExternalArrUB; I++) {
            yield return I;
        }
    }

    // 木のノードのUBを返す
    internal long GetUB()
    {
        return mExternalArrUB;
    }

    // コンストラクタ(外部配列のUBのみ指定)
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

        // フェニック木の外部配列は0オリジンで、
        // フェニック木の内部配列は1オリジンなため、2を足す
        mBitArr = new long[pExternalArrUB + 2];
    }

    // コンストラクタ(初期化用の配列指定)
    internal Fenwick_Tree(long[] pArr)
        : this(pArr.GetUpperBound(0))
    {
        for (long I = 0; I <= pArr.GetUpperBound(0); I++) {
            this.Add(I, pArr[I]);
        }
    }

    // コンストラクタ(初期化用のList指定)
    internal Fenwick_Tree(List<long> pList)
        : this(pList.Count - 1)
    {
        for (int I = 0; I <= pList.Count - 1; I++) {
            this.Add(I, pList[I]);
        }
    }

    // Indのチェック
    private void IndCheck(long pInd)
    {
        if (pInd < 0) throw new Exception("pInd < 0");
        if (mExternalArrUB < pInd) throw new Exception("UB < pInd");
    }

    // Indの大小チェック
    private void IndRangeCheck(long pSta, long pEnd)
    {
        IndCheck(pSta);
        IndCheck(pEnd);
        if (pSta > pEnd) throw new Exception("pSta > pEnd");
    }

    // インデクサ
    internal long this[long pInd]
    {
        get { return GetSum(pInd, pInd); }
        set { Add(pInd, value - GetSum(pInd, pInd)); }
    }

    // [pSta,pEnd] のSumを返す
    internal long GetSum(long pSta, long pEnd)
    {
        IndRangeCheck(pSta, pEnd);

        long Result = GetSum(pEnd);
        if (pSta > 0) {
            Result -= GetSum(pSta - 1);
        }
        return Result;
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        IndCheck(pEnd);

        pEnd++; // 1オリジンに変更

        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        IndCheck(pI);

        pI++; // 1オリジンに変更

        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

解説

ガリガリ君のアタリとハズレを意識します。

最初に何個かガリガリ君を買う。
その後は、アタリ棒のみで、ガリガリ君をもらえるだけ貰う。

と考えると、最初に×の累積和を求めておく。
最初に買った個数での
○の区間和が、アタリ棒の数。
×の区間和が、×の累積和と参照する際の調整値。
として、
IndのList[累積和]で二分探索で解けます。