LeetCode/cpp/greatest-common-divisor-traversal.cpp

#include <algorithm>
#include <cmath>
#include <unordered_map>
#include <vector>

namespace {
struct factors_iterator {
    factors_iterator() : f(-1) {}
    factors_iterator(int f, int n) : f(f), n(n), upper(std::sqrt(n)) {}

    auto operator!=(const factors_iterator &other) const -> bool {
        return f != other.f;
    }

    auto operator*() const -> int { return f; }

    auto operator++() -> factors_iterator & {
        // ‹f› yielded a prime
        if (f > upper || n == 1) {
            f = -1;
            return *this;
        }

        do {
            ++f;
        } while (f <= upper && n % f != 0);

        // ‹n› is a prime
        if (f > upper) {
            f = n;
        }

        while (n % f == 0) {
            n /= f;
        }

        return *this;
    }

  private:
    int f;
    int n;
    int upper;
};

struct factors {
    factors(int n) : n(n) {}
    auto begin() const -> factors_iterator { return ++factors_iterator(1, n); }
    auto end() const -> factors_iterator { return factors_iterator(); }

  private:
    int n;
};

void dfs(std::unordered_map<int, std::vector<int>> &graph,
         std::vector<bool> &visited, int u) {
    visited[u] = true;
    for (auto v : graph[u]) {
        if (!visited[v]) {
            dfs(graph, visited, v);
        }
    }
}
} // namespace

class Solution {
  public:
    bool canTraverseAllPairs(const std::vector<int> &nums) {
        auto n = nums.size();

        std::unordered_map<int, int> gcd;
        std::unordered_map<int, std::vector<int>> neighbors;

        for (auto i = 0u; i < n; ++i) {
            for (auto f : factors(nums[i])) {
                if (gcd.contains(f)) {
                    int u = i;
                    int v = gcd[f];
                    neighbors[u].push_back(v);
                    neighbors[v].push_back(u);
                }

                gcd[f] = i;
            }
        }

        std::vector<bool> visited(n, false);
        dfs(neighbors, visited, 0);

        return std::find(visited.begin(), visited.end(), false) ==
               visited.end();
    }
};

#ifdef _MF_TEST
#include <gtest/gtest.h>

TEST(examples, no_1) {
    Solution s;
    EXPECT_TRUE(s.canTraverseAllPairs(std::vector{2, 3, 6}));
}
TEST(examples, no_2) {
    Solution s;
    EXPECT_FALSE(s.canTraverseAllPairs(std::vector{3, 9, 5}));
}
TEST(examples, no_3) {
    Solution s;
    EXPECT_TRUE(s.canTraverseAllPairs(std::vector{4, 3, 12, 8}));
}

TEST(factors, prime_2) {
    std::vector<int> fs;
    for (auto f : factors(2)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, std::vector{2});
}
TEST(factors, prime_7) {
    std::vector<int> fs;
    for (auto f : factors(7)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, std::vector{7});
}
TEST(factors, prime_29) {
    std::vector<int> fs;
    for (auto f : factors(29)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, std::vector{29});
}
TEST(factors, composite_6) {
    std::vector<int> fs;
    for (auto f : factors(6)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, (std::vector{2, 3}));
}
TEST(factors, composite_15) {
    std::vector<int> fs;
    for (auto f : factors(15)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, (std::vector{3, 5}));
}
TEST(factors, composite_1024) {
    std::vector<int> fs;
    for (auto f : factors(1024)) {
        fs.push_back(f);
    }
    EXPECT_EQ(fs, std::vector{2});
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}
#endif