Doctor and Patient
unknown
c_cpp
2 years ago
1.7 kB
10
Indexable
bool dfs(const vector<int> &a, const vector<int> &b, const vector<unordered_set<int>>& con, vector<int> &c, int i, int color, vector<int> &all) { if (c[i] == color) { return true; } if (c[i] >= 0 && c[i] != color) { return false; } c[i] = color; all.push_back(i); for (int x : con[color]) { if (x != i && !dfs(a, b, con, c, x, color == (a[x] - 1) ? (b[x] - 1) : (a[x] - 1), all)) { return false; } } return true; } int solution(const vector<int> &a, const vector<int> &b, int s) { vector<unordered_set<int>> con(s); const int n = a.size(); for (int i = 0; i < n; ++i) { con[a[i] - 1].insert(i); con[b[i] - 1].insert(i); } vector<int> c(n, -1); for (int i = 0; i < n; ++i) { if (c[i] < 0) { vector<int> all; if (!dfs(a, b, con, c, i, a[i] - 1, all)) { while (!all.empty()) { c[all.back()] = -1; all.pop_back(); } if (!dfs(a, b, con, c, i, b[i] - 1, all)) { return 0; } } for (int x : all) { con[a[x] - 1].erase(x); con[b[x] - 1].erase(x); } } } return 1; } int main() { cout << solution({1, 1, 3}, {2, 2, 1}, 3) << endl; cout << solution({3, 2, 3, 1}, {1, 3, 1, 2}, 3) << endl; cout << solution({2, 5, 6, 5}, {5, 4, 2, 2}, 8) << endl; cout << solution({1, 2, 1, 6, 8, 7, 8}, {2, 3, 4, 7, 7, 8, 7}, 10) << endl; return 0; }
Editor is loading...