Doctor and Patient

 avatar
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...