Short and O(n), Python and C++


  • 22

    I love this problem. Came up with a linear time solution almost 13 years ago and now I still correctly remembered that it was problem 136 at the now defunct acm.uva.es :-)


    Python ... Simple Precompute ... 64 ms

    It's fastest to precompute and store all possibilities for lookup, and it's simplest to just generate them out of order and then sort them.

    class Solution:
        ugly = sorted(2**a * 3**b * 5**c
                      for a in range(32) for b in range(20) for c in range(14))
        def nthUglyNumber(self, n):
            return self.ugly[n-1]
    

    Python ... O(n), generate first n in order ... 308 ms

    My version of epsilon0's solution. It's nicer than my own old one.

    This generates the first n ugly numbers, in order from smallest to largest, in O(n) time. For each prime 2, 3 and 5, have an index to the next number that can be multiplied with the prime to produce a new ugly number. Update the three indexes and then add the smallest of the three candidate ugly numbers.

    def nthUglyNumber(self, n):
        ugly = [1]
        i2 = i3 = i5 = 0
        while len(ugly) < n:
            while ugly[i2] * 2 <= ugly[-1]: i2 += 1
            while ugly[i3] * 3 <= ugly[-1]: i3 += 1
            while ugly[i5] * 5 <= ugly[-1]: i5 += 1
            ugly.append(min(ugly[i2] * 2, ugly[i3] * 3, ugly[i5] * 5))
        return ugly[-1]
    

    Python ... O(n), generate first n in order with heapq.merge ... 416 ms

    I like using heapq.merge.

    def nthUglyNumber(self, n):
        q2, q3, q5 = [2], [3], [5]
        ugly = 1
        for u in heapq.merge(q2, q3, q5):
            if n == 1:
                return ugly
            if u > ugly:
                ugly = u
                n -= 1
                q2 += 2 * u,
                q3 += 3 * u,
                q5 += 5 * u,
    

    Works, but might not be totally safe, both because I extend lists while iterating over them and because at the start, q2 becomes empty for a moment, which I think allows heapq.merge to drop it (at least in CPython it just doesn't, as it first yields the value before trying to get the next one from that list, and when it continues, the list isn't empty anymore).


    C++ ... O(n), generate first n in order ... 12 ms

    C++ version of the second Python solution, though adding extra variables for the three candidates to prevent re-multiplication.

    int nthUglyNumber(int n) {
        vector<int> ugly(n, 1);
        int c2 = 2, c3 = 3, c5 = 5;
        int i2 = 0, i3 = 0, i5 = 0;
        for (int i=1; i<n; ++i) {
            int last = ugly[i] = min(c2, min(c3, c5));
            while (c2 <= last) c2 = 2 * ugly[++i2];
            while (c3 <= last) c3 = 3 * ugly[++i3];
            while (c5 <= last) c5 = 5 * ugly[++i5];
        }
        return ugly[n-1];
    }
    

    C++ ... O(n), generate on demand and remember ... 4 ms

    Same as the previous, but I keep everything in static variables and only compute more when needed. It's faster, and I think pretty neat.

    int nthUglyNumber(int n) {
        static vector<int> ugly {1};
        static int last(1);
        static int c2(2), c3(3), c5(5);
        static int i2(0), i3(0), i5(0);
        while (ugly.size() < n) {
            while (c2 <= last) c2 = 2 * ugly[++i2];
            while (c3 <= last) c3 = 3 * ugly[++i3];
            while (c5 <= last) c5 = 5 * ugly[++i5];
            ugly.push_back(last = min(c2, min(c3, c5)));
        }
        return ugly[n-1];
    }
    

    C++ ... Simple Precompute ... 4 ms

    Precompute all possibilities in easy order and sort them.

    int nthUglyNumber(int n) {
        static vector<int> ugly;
        long long a, b, c, m = INT_MAX;
        if (ugly.empty()) {
            for (a=1; a<=m; a*=2)
                for (b=a; b<=m; b*=3)
                    for (c=b; c<=m; c*=5)
                        ugly.push_back(c);
            sort(begin(ugly), end(ugly));
        }
        return ugly[n-1];
    }
    

    My own old one

    Python version of my own O(n) C++ solution from almost 13 years ago. Good old times...
    Maybe I'll try to improve it, but mainly I wanted to see how a pretty much direct translation would look.

    def nthUglyNumber(self, n):
        factor = 2, 3, 5
        lists = [collections.deque([1]) for _ in range(3)]
        for _ in range(n - 1):
            next = [lists[i][0] * factor[i] for i in range(3)]
            winner = min(range(3), key=lambda j: next[j])
            for i in range(winner, 3):
                lists[i] += next[winner],
            lists[winner].popleft()
        return lists[2][-1]

  • 0

    Hi, Stefan. Could you please explain "compute more when needed" in the second C++ solution? I did not get it. And it seems that using static has some strict restrictions? I try to make all the variables in the first C++ solution to be static and it met "wrong answer". The problematic code is as follows.

    class Solution {
    public:
        int nthUglyNumber(int n) {
            static vector<int> ugly(n, 1);
            static int c2 = 2, c3 = 3, c5 = 5;
            static int i2 = 0, i3 = 0, i5 = 0;
            for (int i=1; i<n; ++i) {
                static int last = ugly[i] = min(c2, min(c3, c5));
                while (c2 <= last) c2 = 2 * ugly[++i2];
                while (c3 <= last) c3 = 3 * ugly[++i3];
                while (c5 <= last) c5 = 5 * ugly[++i5];
            }
            return ugly[n-1];
        }
    }; 
    

    BTW, just didn't realize the precomputing solution and it is amazingly fast!


  • 0

    Static variables aren't lost between calls, their data is kept and they're also not initialized twice. So if I'm called with n = 90 and my ugly already has 70 numbers due to previous calls, then I just have to compute and append 20 more. And if I then get called with n = 80, I don't have to do any work, as I already have 90. That's why my loop goes while (ugly.size() < n) and I use push_back.

    Your solution has two problems:

    • Your ugly always has the size of the first n you ever got. When you get larger n in later calls, you're writing behind the end of the vector, which is dangerous.
    • Your i always restarts from 1, but the static variables are already ahead. Makes no sense, really.

  • 0

    Hi, Stefan. Thank you! Well, so you mean in the above function, if I call Solution().nthUglyNumber(2) then ugly will be of length 2 and then even if I call Solution().nthUglyNumber(3), ugly will still be of length 2?

    BTW, is static int c2(2); the same as static int c2 = 2;? Thanks!


  • 0

    Yes, the length will always stay the same, as you initialize it only once and never change it.

    Not sure about c2(2). I just tried that because c2 = 2 didn't work and I wasn't sure about it, but the bug was something else and then I was too lazy to change it back :-)

    I overlooked that last also shouldn't be static. I "minimally-fixed" that code now so it gets accepted, but I really prefer the while-loop and push_back.

    int nthUglyNumber(int n) {
        static vector<int> ugly(2000, 1);
        static int c2 = 2, c3 = 3, c5 = 5;
        static int i2 = 0, i3 = 0, i5 = 0;
        static int i = 1;
        for (; i<n; ++i) {
            int last = ugly[i] = min(c2, min(c3, c5));
            while (c2 <= last) c2 = 2 * ugly[++i2];
            while (c3 <= last) c3 = 3 * ugly[++i3];
            while (c5 <= last) c5 = 5 * ugly[++i5];
        }
        return ugly[n-1];
    }
    

  • 0

    Hi, Stefan. Thanks! I get some sense of static now. Well, learned it during my freshman year. Though just 6 years ago, but almost totally forget it...


  • 0
    H

    Using "static" variables to store results among test cases is really smart, but somehow I feel like it's cheating.
    And when I see 4ms for precompute one, I realize INT_MAX is not a big number.

    Your answers are opening and inspire me a lot.


  • 0
    J

    Hello Stefan, your work really inspired me, I had never thought how LeetCode test our code, so at first I feel a little confused when I see "static" and the difference of runtime.And now I know, Thank you;-)

    BTW,I think you can also use "if" instead of "while" which I think is more strctly logical.There can't be more than one loop to adjust l1 or l2 or l3.

    class Solution {
    public:
        int nthUglyNumber(int n) {      
            static int l1 = 1, l2 = 1, l3 = 1;
            static int i1 = 0, i2 = 0, i3 = 0;
            static vector<int> res{1};
            while(res.size() < n){
                int u = min(min(l1 * 2, l2 * 3), l3 * 5);
                res.push_back(u);
                if(l1 * 2 <= u)
                    l1 = res[++i1];
                if(l2 * 3 <= u)
                    l2 = res[++i2];
                if(l3 * 5 <= u)
                    l3 = res[++i3];
            }
            return res[n - 1];
        }
    };
    

  • 0
    T

    Hi, Stefan, in Python ... Simple Precompute ... 64 ms solution, how did you come up with the numbers 32, 20 14? And len(ugly) == 8960, the auther said that n does not exceed 1690, so the numbers should be ? Thanks a lot.


  • 0

    @Thomas_Y I don't remember how, probably I either computed them with logarithms (like log5(232) ≈ 13.8) or I found them experimentally (like max(e for e in range(32) if 5**e < 2**32)). And yes, I do compute more/larger numbers for simplicity. Doesn't hurt, does it?

    For even more simplicity I could've just used 32 (or 31) for all three. Could've also made it shorter by using for a, b, c in itertools.product(range(31), repeat=3)).


  • 0
    T

    @StefanPochmann I got it, thanks very much. You inspired me a lot.


Log in to reply
 

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