DSL_2_C: Range Search (kD Tree)

問題へのリンク

C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

// Q044 領域探索 https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_C&lang=jp
class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("6");
            WillReturn.Add("2 1");
            WillReturn.Add("2 2");
            WillReturn.Add("4 2");
            WillReturn.Add("6 2");
            WillReturn.Add("3 3");
            WillReturn.Add("5 4");
            WillReturn.Add("2");
            WillReturn.Add("2 4 0 4");
            WillReturn.Add("4 10 2 5");
            //0
            //1
            //2
            //4
            //
            //2
            //3
            //5
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

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

        int N = int.Parse(InputList[0]);

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

        int CurrPointNo = 0;
        var PointInfoList = new List<PointInfoDef>();
        foreach (string EachStr in InputList.Skip(1).Take(N)) {
            SplitAct(EachStr);
            PointInfoDef WillAdd;
            WillAdd.PointNo = CurrPointNo++;
            WillAdd.X = wkArr[0];
            WillAdd.Y = wkArr[1];
            PointInfoList.Add(WillAdd);
        }

        var InsKdTree = new KdTree(PointInfoList.ToArray());
        // InsKdTree.DebugPrint();

        foreach (string EachStr in InputList.Skip(1 + N + 1)) {
            SplitAct(EachStr);
            InsKdTree.mPointList.Clear();
            InsKdTree.Find(0, wkArr[0], wkArr[1], wkArr[2], wkArr[3], 0);

            InsKdTree.mPointList.Sort();
            var sb = new System.Text.StringBuilder();
            InsKdTree.mPointList.ForEach(pX => sb.Append(pX + Environment.NewLine));
            Console.WriteLine(sb.ToString());
        }
    }
}

internal struct PointInfoDef
{
    internal int PointNo;
    internal int X;
    internal int Y;
}

//KdTreeクラス
internal class KdTree
{
    const int NIL = -1;
    private int mCurrNode;
    private PointInfoDef[] mArr;

    // コンストラクタ
    internal KdTree(PointInfoDef[] pArr)
    {
        mCurrNode = 0;
        mArr = pArr;
        Make2DTree(0, mArr.GetUpperBound(0), 0);
    }

    private struct NodeInfoDef
    {
        internal int ArrInd;
        internal int LeftNode;
        internal int RightNode;
    }
    private Dictionary<int, NodeInfoDef> mNodeInfoDict = new Dictionary<int, NodeInfoDef>();

    private class SortX : System.Collections.Generic.IComparer<PointInfoDef>
    {
        public int Compare(PointInfoDef a, PointInfoDef b)
        {
            return a.X.CompareTo(b.X);
        }
    }

    private class SortY : System.Collections.Generic.IComparer<PointInfoDef>
    {
        public int Compare(PointInfoDef a, PointInfoDef b)
        {
            return a.Y.CompareTo(b.Y);
        }
    }

    private int Make2DTree(int pLeftNode, int pRightNode, int pDepth)
    {
        // 部分ソート (pLeftNodeからpRightNodeまで)
        if (pDepth % 2 == 0) {
            Array.Sort(mArr, pLeftNode, pRightNode - pLeftNode + 1, new SortX());
        }
        else {
            Array.Sort(mArr, pLeftNode, pRightNode - pLeftNode + 1, new SortY());
        }

        int Mid = (pLeftNode + pRightNode) / 2;

        int Node = mCurrNode++;
        NodeInfoDef WillAdd;
        WillAdd.ArrInd = Mid;

        // 左部分木の設定
        WillAdd.LeftNode = NIL;
        if (pLeftNode <= Mid - 1) {
            WillAdd.LeftNode = Make2DTree(pLeftNode, Mid - 1, pDepth + 1);
        }

        // 右部分木の設定
        WillAdd.RightNode = NIL;
        if (Mid + 1 <= pRightNode) {
            WillAdd.RightNode = Make2DTree(Mid + 1, pRightNode, pDepth + 1);
        }

        mNodeInfoDict[Node] = WillAdd;

        return Node;
    }

    internal List<int> mPointList = new List<int>();

    internal void Find(int pCurrNode, int pXFrom, int pXTo, int pYFrom, int pYTo, int pDepth)
    {
        NodeInfoDef CurrNodeInfo = mNodeInfoDict[pCurrNode];
        int CurrX = mArr[CurrNodeInfo.ArrInd].X;
        int CurrY = mArr[CurrNodeInfo.ArrInd].Y;

        if (pXFrom <= CurrX && CurrX <= pXTo
         && pYFrom <= CurrY && CurrY <= pYTo) {
            mPointList.Add(mArr[CurrNodeInfo.ArrInd].PointNo);
        }

        if (pDepth % 2 == 0) {
            // 左部分木を探索
            if (CurrNodeInfo.LeftNode != NIL && pXFrom <= CurrX) {
                Find(CurrNodeInfo.LeftNode, pXFrom, pXTo, pYFrom, pYTo, pDepth + 1);
            }

            // 右部分木を探索
            if (CurrNodeInfo.RightNode != NIL && CurrX <= pXTo) {
                Find(CurrNodeInfo.RightNode, pXFrom, pXTo, pYFrom, pYTo, pDepth + 1);
            }
        }
        else {
            // 左部分木を探索
            if (CurrNodeInfo.LeftNode != NIL && pYFrom <= CurrY) {
                Find(CurrNodeInfo.LeftNode, pXFrom, pXTo, pYFrom, pYTo, pDepth + 1);
            }

            // 右部分木を探索
            if (CurrNodeInfo.RightNode != NIL && CurrY <= pYTo) {
                Find(CurrNodeInfo.RightNode, pXFrom, pXTo, pYFrom, pYTo, pDepth + 1);
            }
        }
    }

    internal void DebugPrint()
    {
        foreach (var EachPair in mNodeInfoDict.OrderBy(pX => pX.Key)) {
            Console.WriteLine("mTreeInfoDict[{0}] ArrInd={1} LeftNode={2} RightNode={3}",
                EachPair.Key,
                EachPair.Value.ArrInd,
                EachPair.Value.LeftNode,
                EachPair.Value.RightNode);
        }
    }
}

解説

KdTreeクラスを使用してます。