ABC255-D ±1 Operation 2

問題へのリンク

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 3");
            WillReturn.Add("6 11 2 5 5");
            WillReturn.Add("5");
            WillReturn.Add("20");
            WillReturn.Add("0");
            //10
            //71
            //29
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("10 5");
            WillReturn.Add("1000000000 314159265 271828182 141421356 161803398 0 777777777 255255255 536870912 998244353");
            WillReturn.Add("555555555");
            WillReturn.Add("321654987");
            WillReturn.Add("1000000000");
            WillReturn.Add("789456123");
            WillReturn.Add("0");
            //3316905982
            //2811735560
            //5542639502
            //4275864946
            //4457360498
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] AArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();

        Array.Sort(AArr);
        long UB = AArr.GetUpperBound(0);
        long[] RunSumSei = (long[])AArr.Clone();
        long[] RunSumRev = (long[])AArr.Clone();

        for (long I = 1; I <= UB; I++) {
            RunSumSei[I] += RunSumSei[I - 1];
        }
        for (long I = UB - 1; 0 <= I; I--) {
            RunSumRev[I] += RunSumRev[I + 1];
        }

        var sb = new System.Text.StringBuilder();
        foreach (string EachStr in InputList.Skip(2)) {
            long X = long.Parse(EachStr);

            long Ind1 = ExecNibunhou_LowerOrEqual_Max(X, AArr);
            long Ind2 = ExecNibunhou_UpperBound(X, AArr);

            long LeftSum = 0;
            if (Ind1 > -1) {
                long Cnt = Ind1 + 1;
                LeftSum += Math.Abs(X * Cnt - RunSumSei[Ind1]);
            }
            long RightSum = 0;
            if (Ind2 > -1) {
                long Cnt = UB - Ind2 + 1;
                LeftSum += Math.Abs(X * Cnt - RunSumRev[Ind2]);
            }
            sb.Append(LeftSum + RightSum);
            sb.AppendLine();
        }
        Console.Write(sb.ToString());
    }

    // 二分法で、Val以下で最大の値を持つ、添字を返す
    static int ExecNibunhou_LowerOrEqual_Max(long pVal, long[] pArr)
    {
        // 最後の要素がVal以下の特殊ケース
        if (pVal >= pArr.Last()) {
            return pArr.GetUpperBound(0);
        }
        // 最初の要素がVal超えの特殊ケース
        if (pVal < pArr[0]) {
            return -1;
        }

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

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

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

    // 二分法で、Val超えで最小の値を持つ、添字を返す
    static int ExecNibunhou_UpperBound(long pVal, long[] pArr)
    {
        // 最後の要素が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;
    }
}

解説

6 9 2 5 8 9
という配列で、クエリでXに7が来た時を考えます。

配列はソートしても問題ないので、ソートします。
2 5 6 8 9 9

考察すると、
前後に累積和を持っておいて、
X以下で最大の添字と
X超えで最小の添字を
二分探索すれば良いと分かります。