Official

D - Magical Cookies Editorial by en_translator


Basically, perform the procedure just as instructed in the problem statement; however, you must not spend \(\Theta(H)\) or \(\Theta(W)\) to determine if the cookies in a row or a column have the same color, or the steps 1. and 2. in an iteration costs \(\Theta(HW)\) time, for a total complexity of \(\Theta(HW(H+W))\) time in the worst case, which leads to TLE (Time Limit Exceeded).

Let us try to optimize it based on the constraints that there are only a few kinds of colors. Suppose that there are \(\sigma\) kinds of colors. In this problem, \(\sigma = 26\). For each row, one can manage the number of cookies with colors a, b, \(\ldots\), and z while performing the operation, so that one can check if the colors of the cookies in a row or column is all the same in an \(O(\sigma)\) time. Therefore, that decision problem can be solved in a total of \(O(\sigma HW)\) time. Marking and removing cookies can be done in a total of \(O(HW)\) time, because such operations are performed only once for each cookie. Therefore, the problem has been solved in a total time complexity of \(O(\sigma HW)\).

Sample code
Note that we do not actually mark cookies, but manage the indices of the rows and columns to mark instead.

#include <iostream>
#include <vector>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;

int main() {
	int h, w;
	cin >> h >> w;
	vector<vector<char>> c(h, vector<char>(w));
	rep(i, h) rep(j, w) cin >> c[i][j];
	vector<vector<int>> x(h, vector<int>(26));
	vector<vector<int>> y(w, vector<int>(26));
	rep(i, h) rep(j, w) {
		x[i][c[i][j] - 'a']++;
		y[j][c[i][j] - 'a']++;
	}
	int hc = h, wc = w;
	vector<bool> fx(h), fy(w);
	rep(_, h + w) {
		vector<pair<int, int>> ux, uy;
		rep(i, h) {
			if (fx[i]) continue;
			rep(j, 26) {
				if (x[i][j] == wc && wc >= 2) {
					ux.push_back({ i, j });
				}
			}
		}
		rep(i, w) {
			if (fy[i]) continue;
			rep(j, 26) {
				if (y[i][j] == hc && hc >= 2) {
					uy.push_back({ i, j });
				}
			}
		}
		for (pair<int, int> p : ux) {
			fx[p.first] = true;
			rep(i, w) y[i][p.second]--;
			hc--;
		}
		for (pair<int, int> p : uy) {
			fy[p.first] = true;
			rep(i, h) x[i][p.second]--;
			wc--;
		}
	}
	cout << hc * wc << '\n';
}

posted:
last update: