algovault/code_library/data_structures/potential_dsu.cpp at master · sgtlaugh/algovault · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
/***
 *
 * DSU with Potentials
 * Union-find that also knows how much each value differs from the others in its set
 *
 * Complexity: amortized O(α(n)) per operation
 *
 * PotentialDSU dsu(n, mod = 0): nodes 0..n, all values unknown
 * dsu.relate(a, b, d): record x[b] = x[a] + d, false and no change if that contradicts what is known
 * dsu.diff(a, b): x[b] - x[a], a and b must be in the same set
 * dsu.find(a): the root of a's set, the same node for every member
 * dsu.same(a, b): true if a and b are in the same set, so diff(a, b) is known
 * dsu.size(a): number of nodes in a's set
 *
 * Notes:
 *   - mod = 0 works on plain integers, every implied difference must fit in long long
 *   - mod > 0 works modulo mod, mod <= 1e18
 *   - mod = 2 is parity, relate(a, b, 1) means a and b differ
 *   - Exact ratios x[b] / x[a] = p / q: one instance per prime, d = that prime's exponent in p minus its exponent in q
 *
***/

#include <bits/stdtr1c++.h>

using namespace std;

struct PotentialDSU{
    long long mod;
    vector<int> parent, sz;
    vector<long long> up;  /// x[i] = x[parent[i]] + up[i], so x[i] = x[root] + up[i] right after find(i)

    PotentialDSU(int n, long long mod = 0) : mod(mod), parent(n + 1), sz(n + 1, 1), up(n + 1, 0){
        assert(0 <= mod && mod <= (long long)1e18);
        iota(parent.begin(), parent.end(), 0);
    }

    bool relate(int a, int b, long long d){
        int ra = find(a), rb = find(b);

        d = norm(d);
        if (ra == rb) return norm(up[b] - up[a]) == d;

        /// d + up[a] first: it is the implied difference x[b] - x[ra], so it fits
        long long w = norm(d + up[a] - up[b]);
        if (sz[ra] < sz[rb]) swap(ra, rb), w = norm(-w);
        parent[rb] = ra, up[rb] = w, sz[ra] += sz[rb];
        return true;
    }

    long long diff(int a, int b){
        int ra = find(a), rb = find(b);
        assert(ra == rb);
        return norm(up[b] - up[a]);
    }

    int find(int i){
        if (parent[i] == i) return i;
        int root = find(parent[i]);
        up[i] = norm(up[i] + up[parent[i]]);
        return parent[i] = root;
    }

    bool same(int a, int b){
        return find(a) == find(b);
    }

    int size(int a){
        return sz[find(a)];
    }

    long long norm(long long x) const{
        if (!mod) return x;
        x %= mod;
        return x < 0 ? x + mod : x;
    }
};

int main(){
    /// Heights: B is 3 taller than A and C is 2 taller than B
    int A = 0, B = 1, C = 2, D = 3;
    PotentialDSU heights(3);
    assert(heights.relate(A, B, 3));
    assert(heights.relate(B, C, 2));
    assert(heights.diff(A, C) == 5);
    assert(heights.diff(C, A) == -5);
    assert(heights.relate(A, C, 5));
    assert(!heights.relate(C, A, 1));  /// A is 5 shorter than C, not 1 taller
    assert(!heights.same(A, D));
    assert(heights.size(A) == 3);

    /// Parity: an odd cycle contradicts
    PotentialDSU parity(3, 2);
    assert(parity.relate(1, 2, 1));
    assert(parity.relate(2, 3, 1));
    assert(!parity.relate(3, 1, 1));
    assert(parity.diff(1, 3) == 0);

    return 0;
}