ABC023-C 収集王

問題へのリンク

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

    static void Main()
    {
        List<string> InputList = GetInputList();

        long[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(pX => long.Parse(pX)).ToArray();

        SplitAct(InputList[0]);
        long UB_X = wkArr[1] - 1;
        long UB_Y = wkArr[0] - 1;
        long K = wkArr[2];

        // Y座標のSet[X座標]なDict
        var AmeYSetDict = new Dictionary<long, HashSet<long>>();

        // 横計[Y座標]な配列
        long[] YokokeiArr = new long[UB_Y + 1];

        // 縦計[X座標]な配列
        long[] TatekeiArr = new long[UB_X + 1];

        foreach (string EachStr in InputList.Skip(2)) {
            SplitAct(EachStr);

            long CurrX = wkArr[1] - 1;
            long CurrY = wkArr[0] - 1;

            if (AmeYSetDict.ContainsKey(CurrX) == false) {
                AmeYSetDict[CurrX] = new HashSet<long>();
            }
            AmeYSetDict[CurrX].Add(CurrY);

            YokokeiArr[CurrY]++;
            TatekeiArr[CurrX]++;
        }

        // Y座標のSet[横計]なDict
        var YokokeiYSetDict = new Dictionary<long, HashSet<long>>();

        for (long I = 0; I <= UB_Y; I++) {
            long CurrYokoKei = YokokeiArr[I];
            if (YokokeiYSetDict.ContainsKey(CurrYokoKei) == false) {
                YokokeiYSetDict[CurrYokoKei] = new HashSet<long>();
            }
            YokokeiYSetDict[CurrYokoKei].Add(I);
        }

        long Answer = 0;
        for (long LoopX = 0; LoopX <= UB_X; LoopX++) {
            long CurrCnt = TatekeiArr[LoopX];

            long NeedCnt1 = K - CurrCnt;  // 必要数(ダブルカウントがない場合)
            long NeedCnt2 = NeedCnt1 + 1; // 必要数(ダブルカウントがある場合)

            long AnswerAdd1 = 0;
            if (YokokeiYSetDict.ContainsKey(NeedCnt1)) {
                AnswerAdd1 = YokokeiYSetDict[NeedCnt1].Count;
                if (AmeYSetDict.ContainsKey(LoopX)) {
                    // これだとTLEになる
                    // HashSet<long> wkYSet = new HashSet<long>(YokokeiYSetDict[NeedCnt1]);
                    // wkYSet.ExceptWith(AmeYSetDict[LoopX]);
                    // AnswerAdd1 = wkYSet.Count;

                    // これだとTLEにならない
                    HashSet<long> wkYSet = new HashSet<long>(AmeYSetDict[LoopX]);
                    wkYSet.IntersectWith(YokokeiYSetDict[NeedCnt1]);
                    AnswerAdd1 -= wkYSet.Count;
                }
            }

            long AnswerAdd2 = 0;
            if (YokokeiYSetDict.ContainsKey(NeedCnt2)) {
                if (AmeYSetDict.ContainsKey(LoopX)) {
                    HashSet<long> wkYSet = new HashSet<long>(AmeYSetDict[LoopX]);
                    wkYSet.IntersectWith(YokokeiYSetDict[NeedCnt2]);
                    AnswerAdd2 = wkYSet.Count;
                }
            }
            Answer += AnswerAdd1 + AnswerAdd2;
        }
        Console.WriteLine(Answer);
    }
}

解説

X座標ごとの飴のあるY座標のSet
X座標ごとの縦計
Y座標ごとの横計
を求めてから、
横計ごとのY座標のSetを求めます。

その後の集計で集合演算を使って高速に集計してます。