AtCoderのABC    次のABCの問題へ    前のABCの問題へ

ABC373-E How to Win the Election


問題へのリンク


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 2 16");
            WillReturn.Add("3 1 4 1 5");
            //2 -1 1 -1 0
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("12 1 570");
            WillReturn.Add("81 62 17 5 5 86 15 7 79 26 6 28");
            //79 89 111 117 117 74 112 116 80 107 117 106
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long mN;
    static long mM;
    static long mTotalCnt;
    static long mRestCnt;

    static long[] mAArr;
    static long[] mOriginalArr;
    static long UB;

    static Fenwick_Tree mInsFenwick_Tree;

    struct RangeInfoDef
    {
        internal long RangeSta;
        internal long RangeEnd;
    }

    // ライバル区間のList
    static Dictionary<long, List<RangeInfoDef>> mRivalRangeListDict =
        new Dictionary<long, List<RangeInfoDef>>();

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = InputList[0].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        mN = wkArr[0];
        mM = wkArr[1];
        mTotalCnt = wkArr[2];

        mAArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        mOriginalArr = (long[])mAArr.Clone();
        Array.Sort(mAArr);
        UB = mAArr.GetUpperBound(0);

        mRestCnt = mTotalCnt - mAArr.Sum();

        mInsFenwick_Tree = new Fenwick_Tree(mAArr);

        if (mN <= mM) {
            var AnswerList1 = new List<long>();
            for (long I = 0; I <= UB; I++) {
                AnswerList1.Add(0);
            }
            Console.WriteLine(LongEnumJoin(" ", AnswerList1));
            return;
        }

        // ライバル区間のListを設定
        for (long I = 0; I <= UB; I++) {
            mRivalRangeListDict[I] = new List<RangeInfoDef>();

            long DefaultRangeSta = UB - mM + 1;
            long DefaultRangeEnd = UB;

            if (I < DefaultRangeSta) {
                RangeInfoDef WillAdd;
                WillAdd.RangeSta = DefaultRangeSta;
                WillAdd.RangeEnd = DefaultRangeEnd;
                mRivalRangeListDict[I].Add(WillAdd);
                continue;
            }
            RangeInfoDef WillAdd1;
            WillAdd1.RangeSta = DefaultRangeSta - 1;
            WillAdd1.RangeEnd = I - 1;
            mRivalRangeListDict[I].Add(WillAdd1);

            RangeInfoDef WillAdd2;
            WillAdd2.RangeSta = I + 1;
            WillAdd2.RangeEnd = UB;
            if (WillAdd2.RangeSta <= WillAdd2.RangeEnd) {
                mRivalRangeListDict[I].Add(WillAdd2);
            }
        }

        var AnswerDict = new Dictionary<long, long>();
        for (long I = 0; I <= UB; I++) {
            // 0で達成可能な場合
            if (CanAchieve(I, 0)) {
                AnswerDict[mAArr[I]] = 0;
                continue;
            }

            // mRestCntでも達成できない場合
            if (CanAchieve(I, mRestCnt) == false) {
                AnswerDict[mAArr[I]] = -1;
                continue;
            }

            // 追加得票数Kを二分探索
            long L = 0;
            long R = mRestCnt;

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

                if (CanAchieve(I, Mid)) {
                    R = Mid;
                }
                else {
                    L = Mid;
                }
            }
            AnswerDict[mAArr[I]] = R;
        }

        var AnswerList2 = new List<long>();
        for (long I = 0; I <= UB; I++) {
            AnswerList2.Add(AnswerDict[mOriginalArr[I]]);
        }
        Console.WriteLine(LongEnumJoin(" ", AnswerList2));
    }

    // 添字と追加得票数Kを引数として、当選確定かを返す
    static bool CanAchieve(long pInd, long pK)
    {
        long GetCnt = mAArr[pInd] + pK;
        long ShikiVal = GetCnt + 1;
        long CurrRestCnt = mRestCnt - pK;

        int ResultInd = ExecNibunhou_LowerBound(ShikiVal, mAArr);

        // 負ける可能性があるライバルの人数と、票数合計を求める
        long Cnt = 0;
        long Sum = 0;
        foreach (RangeInfoDef EachRangeInfo in mRivalRangeListDict[pInd]) {
            long RangeSta = EachRangeInfo.RangeSta;
            long RangeEnd = EachRangeInfo.RangeEnd;
            if (ResultInd > -1 && RangeSta <= ResultInd && ResultInd <= RangeEnd) {
                RangeEnd = ResultInd - 1;
            }
            if (RangeSta <= RangeEnd) {
                Cnt += RangeEnd - RangeSta + 1;
                Sum += mInsFenwick_Tree.GetSum(RangeSta, RangeEnd);
            }
        }

        return ShikiVal * Cnt - Sum > CurrRestCnt;
    }

    // 二分法で、Val以上で最小の値を持つ、添字を返す
    static int ExecNibunhou_LowerBound(long pVal, long[] pArr)
    {
        if (pArr.Length == 0) return -1;

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

        int L = 0;
        int R = pArr.GetUpperBound(0);

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

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

    // セパレータとLong型の列挙を引数として、結合したstringを返す
    static string LongEnumJoin(string pSeparater, IEnumerable<long> pEnum)
    {
        string[] StrArr = Array.ConvertAll(pEnum.ToArray(), pX => pX.ToString());
        return string.Join(pSeparater, StrArr);
    }
}

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

    // [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


解説

M=3として、得票数の昇順に考えます。
1 3 5 7 9 11 13 15 17 得票数
0 1 2 3 4  5  6  7  8 添字

最初に、添え字ごとの負ける候補な、ライバルの区間のListを作成します。
添字0のライバルは、[6,8]
添字1のライバルは、[6,8]
添字2のライバルは、[6,8]
添字3のライバルは、[6,8]
添字4のライバルは、[6,8]
添字5のライバルは、[6,8]
添字6のライバルは、[5,5]と[7,8]
添字7のライバルは、[5,6]と[8,8]
添字8のライバルは、[5,7]

後は、何票獲得で当選確定かを二分探索してます。
二分探索用に、
残った票が全て最悪な分布になっても、当選できるかの判定メソッドを作ってます。