ABC393-D Swap to Gather

問題へのリンク

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("7");
            WillReturn.Add("0101001");
            //3
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("3");
            WillReturn.Add("100");
            //0
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("10");
            WillReturn.Add("0101001001");
            //7
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static char[] mCharArr;
    static long UB;

    static Fenwick_Tree mInsFenwick_TreeOne;

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

        var OneList = new List<long>();
        for (long I = 0; I <= UB; I++) {
            if (mCharArr[I] == '1') {
                OneList.Add(I);
            }
        }
        mInsFenwick_TreeOne = new Fenwick_Tree(UB);
        for (long I = 0; I <= UB; I++) {
            if (mCharArr[I] == '1') {
                mInsFenwick_TreeOne[I] = 1;
            }
        }

        if (OneList.Count == 1) {
            Console.WriteLine(0);
            return;
        }

        var AnswerList = new List<long>();
        if (OneList.Count % 2 == 0) {
            long MidInd1 = OneList[OneList.Count / 2];
            long MidInd2 = OneList[OneList.Count / 2 - 1];

            AnswerList.Add(DeriveCost(MidInd1));
            AnswerList.Add(DeriveCost(MidInd2));
        }
        else {
            long MidInd1 = OneList[OneList.Count / 2];
            AnswerList.Add(DeriveCost(MidInd1));
        }
        Console.WriteLine(AnswerList.Min());
    }

    // 中心を引数として、コストを返す
    static long DeriveCost(long pMidInd)
    {
        long Answer = 0;
        for (long I = 0; I <= UB; I++) {
            if (mCharArr[I] != '1') continue;

            if (I < pMidInd) {
                long RangeSta = I;
                long RangeEnd = pMidInd - 1;
                long PlusCost = RangeEnd - RangeSta;
                long MinusCost = -1 + mInsFenwick_TreeOne.GetSum(RangeSta, RangeEnd);
                Answer += PlusCost - MinusCost;
            }
            if (I > pMidInd) {
                long RangeSta = pMidInd + 1;
                long RangeEnd = I;
                long PlusCost = RangeEnd - RangeSta;
                long MinusCost = -1 + mInsFenwick_TreeOne.GetSum(RangeSta, RangeEnd);
                Answer += PlusCost - MinusCost;
            }
        }
        return Answer;
    }
}

// フェニック木
#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]);
        }
    }

    // インデクサ
    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)
    {
        return GetSum(pEnd) - GetSum(pSta - 1);
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long 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)
    {
        pI++; // 1オリジンに変更

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

解説