Share my O(nk) solution from ugly ii and priority queue solution


  • 1
    S
    private:
    struct Node {
        long val;
        int idx, prim; // index of different primes 
        Node(long _v = 0, int _i = 0, int _p = 0) : val(_v), idx(_i), prim(_p) {}
        inline bool operator<(const Node &x) const {
            return val > x.val; // min heap
        }
    };
    
    public:
    // priority queue solution
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        priority_queue<Node> q;
        vector<long> ans(n + 10, 0);
        ans[0] = 1;
        
        // init the first k elements in pq
        for (int &p : primes) q.push(Node(p, 0, p));
        for (int i = 1; i < n; ++i) {
            Node ret = q.top();
            ans[i] = ret.val;
            while(true) {
                ret = q.top(); q.pop();
                ret.val = ans[++ret.idx] * ret.prim; // increase different prime's index
                q.push(ret);
                if (q.empty() || q.top().val != ans[i]) break;
            }
        }
        
        return ans[n - 1];
    }
    
    int findMin(vector<int> & v) {
        if (v.size() <= 0) return -1;
        int ans = v[0];
        for (int x : v) {
            if (x < ans) ans = x;
        }
        return ans;
    }
    
    // refer from the solution of Ugly-number-ii
    // O(n*k) time, O(k) space
    int nthSuperUglyNumber(int n, vector<int>& primes) {
        vector<int> d(n, 0);
        d[0] = 1;
        
        int k = primes.size();
        vector<int> v(primes.begin(), primes.end());
        vector<int> indexs(k, 0); // indexs for different prime factors
        
        for (int i = 1; i < n; ++i) {
            int minV = findMin(v);
            d[i] = minV;
            for (int j = 0; j < k; ++j) {
                if (v[j] == minV) v[j] = primes[j] * d[++indexs[j]];
            }
        }
        
        return d[n - 1];
    }

  • 0
    N

    Thanks for sharing! In the first solution, maybe there is a typo. vector<long> ans(n + 10, 0); could just set the size of vector to n.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.