3336. Find the Number of Subsequences With Equal GCD

Array
Dynamic Programming
Number Theory
Math

TIME

O(n * V^2)

SPACE

O(V^2)

Intuition

Each index of nums is independently placed into seq1, into seq2, or into neither, three choices. What matters about the two growing subsequences is only their current GCDs, and GCD values never exceed max(nums) <= 200. So a DP over (gcd1, gcd2) captures everything, and we count assignments where both are non-empty and the GCDs match.

Approach

State dp[g1][g2] = number of ways to assign the processed prefix so that seq1 has GCD g1 and seq2 has GCD g2. Use 0 to denote "still empty" (with the convention gcd(0, x) = x).
Initialize dp[0][0] = 1 (both empty).
For each num, every existing state (g1, g2) with count c branches into three:
- skip: (g1, g2) unchanged
- add to seq1: (gcd(g1, num), g2)
- add to seq2: (g1, gcd(g2, num)) Accumulate into a fresh next-layer table (mod 1e9 + 7).
Answer: sum over g >= 1 of dp[g][g] both subsequences non-empty and equal GCD.

Using a fresh next table per element keeps each element's three choices independent (no double-application).

Complexity

Time: O(n · V²) where V = max value (≤ 200)
Space: O(V²), two DP tables

Edge Cases

Excluding g = 0 guarantees both subsequences are non-empty
Repeated values (e.g. all 1s) inflate counts, handled naturally by the per-index branching
Disjointness is automatic: each index picks exactly one of the three destinations
Single element → no valid pair (can't fill both) → 0

Code

from typing import List
from math import gcd


class Solution:
    def subsequencePairCount(self, nums: List[int]) -> int:
        """
        3336. Find the Number of Subsequences With Equal GCD
        Time: O(n * V^2)
        Space: O(V^2)
        """
        MOD = 10**9 + 7
        V = max(nums)

        # dp[g1][g2]: ways so far; 0 means "empty" (gcd(0, x) = x).
        dp = [[0] * (V + 1) for _ in range(V + 1)]
        dp[0][0] = 1

        for num in nums:
            nxt = [row[:] for row in dp]  # start from "skip" choice
            for g1 in range(V + 1):
                row = dp[g1]
                for g2 in range(V + 1):
                    c = row[g2]
                    if c == 0:
                        continue
                    ng1 = gcd(g1, num)  # gcd(0, num) == num
                    ng2 = gcd(g2, num)
                    nxt[ng1][g2] = (nxt[ng1][g2] + c) % MOD  # add to seq1
                    nxt[g1][ng2] = (nxt[g1][ng2] + c) % MOD  # add to seq2
            dp = nxt

        return sum(dp[g][g] for g in range(1, V + 1)) % MOD

/**
 * 3336. Find the Number of Subsequences With Equal GCD
 * Time: O(n * V^2)
 * Space: O(V^2)
 */
class Solution {
    private static final int MOD = 1_000_000_007;

    public int subsequencePairCount(int[] nums) {
        int V = 0;
        for (int x : nums) {
            V = Math.max(V, x);
        }

        // dp[g1][g2]: ways so far; 0 means "empty" (gcd(0, x) = x).
        long[][] dp = new long[V + 1][V + 1];
        dp[0][0] = 1;

        for (int num : nums) {
            long[][] nxt = new long[V + 1][V + 1];
            for (int g1 = 0; g1 <= V; g1++) {
                for (int g2 = 0; g2 <= V; g2++) {
                    nxt[g1][g2] = dp[g1][g2]; // "skip" choice
                }
            }
            for (int g1 = 0; g1 <= V; g1++) {
                for (int g2 = 0; g2 <= V; g2++) {
                    long c = dp[g1][g2];
                    if (c == 0) {
                        continue;
                    }
                    int ng1 = gcd(g1, num); // gcd(0, num) == num
                    int ng2 = gcd(g2, num);
                    nxt[ng1][g2] = (nxt[ng1][g2] + c) % MOD; // add to seq1
                    nxt[g1][ng2] = (nxt[g1][ng2] + c) % MOD; // add to seq2
                }
            }
            dp = nxt;
        }

        long ans = 0;
        for (int g = 1; g <= V; g++) {
            ans = (ans + dp[g][g]) % MOD;
        }
        return (int) ans;
    }

    private int gcd(int a, int b) {
        while (b != 0) {
            int t = a % b;
            a = b;
            b = t;
        }
        return a;
    }
}