第8回PAST L K番目の絶対値

問題へのリンク

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("2 2");
            WillReturn.Add("2 -3");
            //2
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("5 10");
            WillReturn.Add("1000000000 1000000000 1000000000 1000000000 1000000000");
            //3000000000
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("7 22");
            WillReturn.Add("3 -1 4 -1 5 -9 2");
            //5
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long mK;

    static List<long> mRunSumList = new List<long>();

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

        long[] AArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();

        mRunSumList.Add(0);
        long CurrRunSum = 0;

        foreach (long EachA in AArr) {
            CurrRunSum += EachA;
            mRunSumList.Add(CurrRunSum);
        }
        mRunSumList.Sort();

        // 0が達成可能な場合
        if (CanAchieve(0)) {
            Console.WriteLine(0);
            return;
        }

        long L = 0;
        long R = AArr.Sum(pX => Math.Abs(pX));

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

            if (CanAchieve(Mid)) {
                R = Mid;
            }
            else {
                L = Mid;
            }
        }
        Console.WriteLine(R);
    }

    // 絶対値がX以下となる区間が、K個以上あるかを返す
    static bool CanAchieve(long pX)
    {
        long RangeCnt = 0;

        foreach (long EachRunSum in mRunSumList) {
            long CurrRangeCnt = GetListRangeValueCnt.GetCnt(mRunSumList, EachRunSum - pX, EachRunSum + pX);
            RangeCnt += CurrRangeCnt;
            if (RangeCnt - mRunSumList.Count >= 2 * mK) return true;
        }
        return RangeCnt - mRunSumList.Count >= 2 * mK;
    }
}

#region GetListRangeValueCnt
// {昇順にソートされたList,Min,Max}を引数とし、
// Min以上Max以下な値の個数を返す
internal class GetListRangeValueCnt
{
    static internal long GetCnt(List<long> pSortedList, long pMinVal, long pMaxVal)
    {
        if (pMinVal > pMaxVal) {
            throw new Exception("pMinVal > pMaxVal");
        }

        int ResultInd1 = ExecNibunhou_LowerBound(pMinVal, pSortedList);
        int ResultInd2 = ExecNibunhou_LowerOrEqual_Max(pMaxVal, pSortedList);

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

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

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

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

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

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

    // 二分法で、Val以下で最大の値を持つ、添字を返す
    static int ExecNibunhou_LowerOrEqual_Max(long pVal, List<long> pList)
    {
        if (pList.Count == 0) return -1;

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

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

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

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

解説

3 -1 4 -1 5 -9 2 -6 5 -3 5
で区間和の絶対値が5以下 ⇔ -5 <= 区間和 <= 5
を考えます。

区間和を高速に求めたいので、累積和を求めます。
  3 -1 4 -1  5 -9 2 -6 5 -3 5
0 3  2 6  5 10  1 3 -3 2 -1 4

3 - 0
2 - 3
2 - 0
6 - 0
6 - 3
6 - 2
5 - 0
5 - 3
5 - 2
5 - 6
と線形に考えて、AVL木で二分探索したり、ウェーブレット木を使っても、TLEするので
下記の考察を行います。

0 と 他の値との差
3 と 他の値との差
2 と 他の値との差
6 と 他の値との差
5 と 他の値との差
で考えます。
(区間和の絶対値なので、引き算の方向は、どっちでも良い、いいかえると、引き算で交換法則が成り立つ)

すると、順番入れ替えで区間数が2倍になる分と、自分と自分のペアで0になる分を考慮すれば、
後は、累積和のListを事前にソートしておいての二分探索が使えるようになり、
TLEを回避できます。