Problem statement

https://leetcode.com/problems/gcd-sort-of-an-array/

Solution

Quite hard problem for python, you need to use couple of tricks to make it AC.

The idea is to the following steps:

  1. Create graph of connectivity. However we can not allow to compare all pairs of numbers, it is O(n^2). Instead what we can do is the opposite idea: for every number 2, 3, 4, 5, ..., we collect all numbers from nums, and add their indexes to dictionary d. There is a way to factorize all numbers between 1 and M in O(M * log M) time, using eratosthenes sieve. Notice that I do it outside function, that is only once for all calls and it allows to have AC in python.
  2. Create graph of connectivites G. For example imagine that d[2] = [1, 7, 8, 9], it means that nums[1], nums[7], nums[8] and nums[9] are divisible by 2. Then we create connections betwen adjacent elements, no need to create all.
  3. Find connected components in the constructed graph.
  4. Sort each connected components and reconstruct mostly sorted array, check if it is equal to fully sorteda array.

Complexity

It is O(M * log M) for time and O(M) for space.

Code

class Solution:
    M = 100000
    arr = [0]*(M+1)
    for i in range(2, M+1):
        for j in range(i, M+1, i):
            if arr[j] == 0: arr[j] = i

    @lru_cache(None)
    def factorize(self, k):
        primes = []
        while self.arr[k] != 0:
            primes += [self.arr[k]]
            k //= self.arr[k]
        return set(primes)
    
    def gcdSort(self, nums):
        d = defaultdict(set)
        for i, num in enumerate(nums):
            for prime in self.factorize(num):
                d[prime].add(i)

        G = defaultdict(list)
        for x in d:
            for a, b in zip(list(d[x]), list(d[x])[1:]):
                G[a] += [b]
                G[b] += [a]

        visited, ans = set(), {}
        comps, ind, n = defaultdict(list), 0, len(nums)

        def dfs(node, i):
            visited.add(node)
            comps[i].append(node)
            for neib in G[node]:
                if neib in visited: continue
                dfs(neib, i)

        for i in range(n):
            if i not in visited:
                dfs(i, ind)
                ind += 1

        for comp in comps.values():
            for i, l in zip(sorted(comp), sorted(nums[t] for t in comp)):
                ans[i] = l

        return [ans[i] for i in range(n)] == sorted(nums)