Leetcode/3548. Equal Sum Grid Partition II/Solution.cs

public class Solution
{
    struct SumInfo
    {
        public long Left;
        public long Right;
        public long Value;
        // ключ - число, значение - индексы, где оно встречается
        public Dictionary<long, SortedSet<int>> Nums;
    }

    public bool CanPartitionGrid(int[][] grid)
    {
        var height = grid.Length;
        var width = grid[0].Length;

        var colSums = new SumInfo[width];
        var rowSums = new SumInfo[height];
        for (var i = 0; i < height; i++)
        {
            for (var j = 0; j < width; j++)
            {
                var num = grid[i][j];
                colSums[j].Value += num;
                rowSums[i].Value += num;
                rowSums[i].Nums ??= new Dictionary<long, SortedSet<int>>();
                colSums[j].Nums ??= new Dictionary<long, SortedSet<int>>();

                rowSums[i].Nums.TryAdd(num, new SortedSet<int>());
                rowSums[i].Nums[num].Add(j);

                colSums[j].Nums.TryAdd(num, new SortedSet<int>());
                colSums[j].Nums[num].Add(i);
            }
            rowSums[i].Left = rowSums[i].Value;
            if (i > 0)
                rowSums[i].Left += rowSums[i - 1].Left;
        }
        for (var i = height - 1; i >= 0; i--)
        {
            rowSums[i].Right = rowSums[i].Value;
            if (i < height - 1)
                rowSums[i].Right += rowSums[i + 1].Right;
        }
        for (var j = 0; j < width; j++)
        {
            colSums[j].Left = colSums[j].Value;
            if (j > 0)
                colSums[j].Left += colSums[j - 1].Left;
        }
        for (var j = width - 1; j >= 0; j--)
        {
            colSums[j].Right = colSums[j].Value;
            if (j < width - 1)
                colSums[j].Right += colSums[j + 1].Right;
        }
        var answer = false;
        for (var i = 0; i < rowSums.Length - 1 && !answer; i++)
        {
            var row = rowSums[i];
            var nextRow = rowSums[i + 1];
            var diff = Math.Abs(row.Left - nextRow.Right);

            var edgesOnly = colSums.Length == 1;
            var colLen = height;
            if (diff == 0 ||
            (row.Left < nextRow.Right && CanDelete(diff, colSums, i, true, i + 1 == rowSums.Length - 1, edgesOnly, colLen)) ||
            (row.Left > nextRow.Right && CanDelete(diff, colSums, i, false, i == 0, edgesOnly, colLen)))
                answer |= true;
        }
        for (var j = 0; j < colSums.Length - 1 && !answer; j++)
        {
            var col = colSums[j];
            var nextCol = colSums[j + 1];
            var diff = Math.Abs(col.Left - nextCol.Right);
            // тут проход по колонкам, значит поиск элемента для удаления надо выполнять по строкам
            // если j == 0 и отрезаем в левой части, то надо смотреть только первую и последнюю строку
            // если j+1 == colSums.Len-1 и отрезаем в правой части, то надо смотреть только первую и последнюю строку
            var edgesOnly = rowSums.Length == 1;
            var rowLen = width;
            if (diff == 0 ||
            (col.Left < nextCol.Right && CanDelete(diff, rowSums, j, true, j + 1 == colSums.Length - 1, edgesOnly, rowLen)) ||
            (col.Left > nextCol.Right && CanDelete(diff, rowSums, j, false, j == 0, edgesOnly, rowLen)))
                answer |= true;
        }
        return answer;
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="target">Целевое число для поиска</param>
    /// <param name="infos">Что перебирать</param>
    /// <param name="idx">Индекс отсечения (для удаления справа строго больше, для удаления слева включительно)</param>
    /// <param name="deleteAfterIdx">Удалять ли спарва от индекса</param>
    /// <param name="firstAndLastOnly">Посмотреть только первый и последний SumInfo</param>
    /// <param name="edgesOnly">В sumInfo брать только крайние элементы</param>
    /// <param name="len">Количество чисел в SumInfo</param>
    /// <returns></returns>
    private static bool CanDelete(long target, SumInfo[] infos, int idx, bool deleteAfterIdx, bool firstAndLastOnly, bool edgesOnly, int len)
    {
        for (var i = 0; i < infos.Length; i++)
        {
            var nums = infos[i].Nums;

            if (nums.TryGetValue(target, out var indexes))
            {
                var before = false;
                var after = false;
                if (edgesOnly)
                {
                    before = indexes.Contains(0) || indexes.Contains(idx);
                    after = indexes.Contains(len - 1) || indexes.Contains(idx+1);
                }
                else
                {
                    before = indexes.GetViewBetween(int.MinValue, idx).Count > 0;
                    after = indexes.GetViewBetween(idx + 1, int.MaxValue).Count > 0;
                }
                if ((before && !deleteAfterIdx) || (after && deleteAfterIdx))
                    return true;
            }

            // если только границы, перейти сразу к последнему элементу
            if (firstAndLastOnly && i == 0)
                i = Math.Max(i, infos.Length - 2);
        }
        return false;
    }
}