DSU with Weighted Edges

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
struct WeightedDSU {
5
    int n;
6
    vector<int> parent, sz;
7
    vector<long long> pot; // pot[x] = val[x] - val[parent[x]]; for roots pot[root] = 0
8
9
    WeightedDSU(int n = 0) { init(n); }
10
11
    void init(int n_) {
12
        n = n_;
13
        parent.resize(n);
14
        sz.assign(n, 1);
15
        pot.assign(n, 0);
16
        iota(parent.begin(), parent.end(), 0);
17
    }
18
19
    // Find root of x and compress path, updating pot[x] to be relative to root
20
    int find(int x) {
21
        if (parent[x] == x) return x;
22
        int p = parent[x];
23
        int r = find(p);
24
        pot[x] += pot[p]; // accumulate offset to the new parent (the root)
25
        parent[x] = r;
26
        return r;
27
    }
28
29
    // Return val[x] - val[root(x)] after ensuring compression
30
    long long weight_to_root(int x) {
31
        find(x);
32
        return pot[x];
33
    }
34
35
    // Return whether a and b are connected
36
    bool same(int a, int b) { return find(a) == find(b); }
37
38
    // Get val[b] - val[a] if connected; throws if not connected
39
    long long diff(int a, int b) {
40
        int ra = find(a);
41
        int rb = find(b);
42
        if (ra != rb) throw runtime_error("Not connected");
43
        return pot[b] - pot[a];
44
    }
45
46
    // Unite with constraint: val[b] - val[a] = w
47
    // Returns true if successful, false if contradiction detected
48
    bool unite(int a, int b, long long w) {
49
        int ra = find(a);
50
        int rb = find(b);
51
        long long pa = pot[a]; // val[a] - val[ra]
52
        long long pb = pot[b]; // val[b] - val[rb]
53
54
        if (ra == rb) {
55
            // Check consistency: (val[b]-val[a]) must equal w
56
            return (pb - pa) == w;
57
        }
58
        // Union by size: attach smaller under larger
59
        if (sz[ra] < sz[rb]) {
60
            // Attach ra under rb. Need pot[ra] = -w + pb - pa
61
            parent[ra] = rb;
62
            pot[ra] = -w + pb - pa;
63
            sz[rb] += sz[ra];
64
        } else {
65
            // Attach rb under ra. Symmetric formula: pot[rb] = w - pb + pa
66
            parent[rb] = ra;
67
            pot[rb] = w - pb + pa;
68
            sz[ra] += sz[rb];
69
        }
70
        return true;
71
    }
72
};
73
74
int main() {
75
    ios::sync_with_stdio(false);
76
    cin.tie(nullptr);
77
78
    // Demo: 5 nodes (0..4). Add constraints and query differences.
79
    WeightedDSU dsu(5);
80
81
    // val[1] - val[0] = 3
82
    cout << (dsu.unite(0, 1, 3) ? "OK" : "Contradiction") << "\n";
83
84
    // val[2] - val[1] = 4  => val[2] - val[0] should be 7
85
    cout << (dsu.unite(1, 2, 4) ? "OK" : "Contradiction") << "\n";
86
87
    // Query val[2] - val[0]
88
    try {
89
        cout << "diff(0,2) = " << dsu.diff(0, 2) << "\n"; // Expect 7
90
    } catch (...) { cout << "Unknown\n"; }
91
92
    // Add a consistent relation: val[2] - val[0] = 7
93
    cout << (dsu.unite(0, 2, 7) ? "OK" : "Contradiction") << "\n"; // OK
94
95
    // Add a contradictory relation: val[2] - val[0] = 8 (conflicts with 7)
96
    cout << (dsu.unite(0, 2, 8) ? "OK" : "Contradiction") << "\n"; // Contradiction
97
98
    // Unknown difference (different components)
99
    try {
100
        cout << "diff(3,4) = " << dsu.diff(3, 4) << "\n";
101
    } catch (...) {
102
        cout << "Unknown\n"; // not connected yet
103
    }
104
105
    // Connect 3 and 4: val[4] - val[3] = -2
106
    cout << (dsu.unite(3, 4, -2) ? "OK" : "Contradiction") << "\n";
107
    cout << "diff(3,4) = " << dsu.diff(3, 4) << "\n"; // -2
108
109
    return 0;
110
}
111

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
struct WeightedDSU {
5
    int n; vector<int> parent, sz; vector<long long> pot;
6
    WeightedDSU(int n=0){init(n);} 
7
    void init(int n_){n=n_; parent.resize(n); sz.assign(n,1); pot.assign(n,0); iota(parent.begin(), parent.end(), 0);}    
8
    int find(int x){ if(parent[x]==x) return x; int p=parent[x]; int r=find(p); pot[x]+=pot[p]; parent[x]=r; return r; }
9
    bool same(int a,int b){ return find(a)==find(b);}    
10
    long long diff(int a,int b){ int ra=find(a), rb=find(b); if(ra!=rb) throw runtime_error("Unknown"); return pot[b]-pot[a]; }
11
    bool unite(int a,int b,long long w){ int ra=find(a), rb=find(b); long long pa=pot[a], pb=pot[b]; if(ra==rb) return (pb-pa)==w; if(sz[ra]<sz[rb]){ parent[ra]=rb; pot[ra]=-w+pb-pa; sz[rb]+=sz[ra]; } else { parent[rb]=ra; pot[rb]=w-pb+pa; sz[ra]+=sz[rb]; } return true; }
12
};
13
14
int main(){
15
    ios::sync_with_stdio(false); cin.tie(nullptr);
16
17
    // Input format (example):
18
    // n q
19
    // Then q lines of two types:
20
    // 1 a b w  -> add constraint val[b] - val[a] = w
21
    // 2 a b    -> query val[b] - val[a] or print "Unknown" if not connected
22
    // Indices are 1-based in input; we convert to 0-based.
23
24
    int n, q; if(!(cin >> n >> q)) return 0;
25
    WeightedDSU dsu(n);
26
    for(int i=0;i<q;i++){
27
        int t; cin >> t;
28
        if(t==1){
29
            int a,b; long long w; cin >> a >> b >> w; --a; --b;
30
            bool ok = dsu.unite(a,b,w);
31
            if(!ok) cout << "Contradiction\n"; else cout << "OK\n";
32
        } else if(t==2){
33
            int a,b; cin >> a >> b; --a; --b;
34
            if(!dsu.same(a,b)) cout << "Unknown\n";
35
            else cout << dsu.diff(a,b) << "\n";
36
        } else {
37
            // Optional: handle invalid type
38
            cout << "Invalid query type\n";
39
        }
40
    }
41
    return 0;
42
}
43

1
#include <bits/stdc++.h>
2
using namespace std;
3
4
struct Vec2 {
5
    long long x=0, y=0;
6
    Vec2(){} Vec2(long long x_, long long y_):x(x_),y(y_){}
7
    Vec2 operator+(const Vec2& o) const { return Vec2(x+o.x, y+o.y); }
8
    Vec2 operator-(const Vec2& o) const { return Vec2(x-o.x, y-o.y); }
9
    Vec2& operator+=(const Vec2& o){ x+=o.x; y+=o.y; return *this; }
10
};
11
12
struct DSU2D {
13
    int n; vector<int> parent, sz; vector<Vec2> pot; // pot[v] = pos[v] - pos[parent[v]] (vector difference)
14
    DSU2D(int n=0){init(n);} 
15
    void init(int n_){ n=n_; parent.resize(n); sz.assign(n,1); pot.assign(n, Vec2()); iota(parent.begin(), parent.end(), 0);}    
16
    int find(int x){ if(parent[x]==x) return x; int p=parent[x]; int r=find(p); pot[x]+=pot[p]; parent[x]=r; return r; }
17
    bool same(int a,int b){ return find(a)==find(b);}    
18
    Vec2 diff(int a,int b){ int ra=find(a), rb=find(b); if(ra!=rb) throw runtime_error("Unknown"); return Vec2(pot[b].x - pot[a].x, pot[b].y - pot[a].y); }
19
    // Enforce pos[b] - pos[a] = w (as a 2D vector)
20
    bool unite(int a,int b, Vec2 w){ int ra=find(a), rb=find(b); Vec2 pa=pot[a], pb=pot[b]; if(ra==rb){ return (pb.x - pa.x == w.x) && (pb.y - pa.y == w.y); }
21
        if(sz[ra] < sz[rb]){ parent[ra]=rb; // pot[ra] = -w + pb - pa
22
            pot[ra] = Vec2(-w.x + pb.x - pa.x, -w.y + pb.y - pa.y); sz[rb]+=sz[ra]; }
23
        else { parent[rb]=ra; // pot[rb] = w - pb + pa
24
            pot[rb] = Vec2(w.x - pb.x + pa.x, w.y - pb.y + pa.y); sz[ra]+=sz[rb]; }
25
        return true; }
26
};
27
28
int main(){
29
    ios::sync_with_stdio(false); cin.tie(nullptr);
30
31
    // Suppose we know relative positions (vectors) between pairs of robots on a plane.
32
    DSU2D dsu(4); // robots 0..3
33
    // pos[1] - pos[0] = (3, 1)
34
    cout << (dsu.unite(0,1, Vec2(3,1)) ? "OK" : "Contradiction") << "\n";
35
    // pos[2] - pos[1] = (-2, 4)
36
    cout << (dsu.unite(1,2, Vec2(-2,4)) ? "OK" : "Contradiction") << "\n";
37
    // Query pos[2] - pos[0] -> should be (1, 5)
38
    try{
39
        Vec2 d = dsu.diff(0,2);
40
        cout << "diff(0,2) = (" << d.x << "," << d.y << ")\n";
41
    }catch(...){ cout << "Unknown\n"; }
42
43
    // Contradictory relation: claim pos[2] - pos[0] = (2,5)
44
    cout << (dsu.unite(0,2, Vec2(2,5)) ? "OK" : "Contradiction") << "\n";
45
46
    return 0;
47
}
48