1982D - Beauty of the mountains ~~~~~
include <bits/stdc++.h>
include <ext/pb_ds/assoc_container.hpp>
include <ext/pb_ds/tree_policy.hpp>
// #include <atcoder/all>
define ordered_set tree<pair<int, int>, null_type, less<pair<int, int>>, rb_tree_tag, tree_order_statistics_node_update>
// #define int long long
const int mod = 1e9 + 7; const int N = 1e5 + 9; const int inf = 1e9;
using namespace __gnu_pbds; using namespace std; // using namespace atcoder;
void solve() { int n, m, k; cin >> n >> m >> k;
vector<vector<int>> mountain(n, vector<int>(m)); vector<string> mountain_type(n); for (int i = 0; i < n; i++) { for (int j = 0; j < m; j++) { cin >> mountain[i][j]; } } for (int i = 0; i < n; i++) { cin >> mountain_type[i]; } int snowy = 0, not_snowy = 0; for (int i = 0; i < n; i++) { for (int j = 0; j < m; j++) { if (mountain_type[i][j] == '0') { not_snowy += mountain[i][j]; } else { snowy += mountain[i][j]; } } } int diff = abs(snowy - not_snowy); vector<vector<int>> prefix(n + 1, vector<int>(m + 1, 0)); for (int i = 1; i <= n; i++) { for (int j = 1; j <= m; j++) { prefix[i][j] = (mountain_type[i - 1][j - 1] - '0') + prefix[i - 1][j] + prefix[i][j - 1] - prefix[i - 1][j - 1]; } } set<int> factors; for (int i = k; i <= n; i++) { for (int j = k; j <= m; j++) { int total_snowy = prefix[i][j] - prefix[i - k][j] - prefix[i][j - k] + prefix[i - k][j - k]; factors.insert(abs(k * k - 2 * total_snowy)); } } int tot_gcd = (factors.size()==0) ? 0 : *factors.begin(); if(factors.size()==0) { if(diff == 0) { cout << "YES" << endl; } else { cout << "NO" << endl; } return; } for (auto x : factors) { if (x == 0) { continue; } tot_gcd = __gcd(tot_gcd, x); } if (diff % tot_gcd == 0) { cout << "YES" << endl; } else { cout << "NO" << endl; }
}
signed main() { ios::sync_with_stdio(0); cin.tie(0); int t = 1; cin >> t; while (t--) { solve(); } return 0; } ~~~~~