Scribbling

[C++] LeetCode 238. Product of Array Except Self 본문

Computer Science/C++

[C++] LeetCode 238. Product of Array Except Self

focalpoint 2023. 8. 18. 06:09

https://leetcode.com/problems/product-of-array-except-self/

class Solution {
public:
    vector<int> productExceptSelf(vector<int>& nums) {
        vector<int> zero_idx;
        int zero_cnt = 0;

        vector<int> dp1(nums.size());
        vector<int> dp2(nums.size());

        int prod = 1;
        for (int i = 0; i < nums.size(); i++) {
            int num = nums[i];
            if (num == 0) {
                zero_idx.push_back(i);
                zero_cnt++;
            }
            else {
                prod *= num;
                dp1[i] = prod;
            }
        }
        prod = 1;
        for (int i = nums.size() - 1; i >= 0; i--) {
            int num = nums[i];
            if (num != 0) {
                prod *= num;
                dp2[i] = prod;
            }
        }
        if (zero_cnt >= 2) {
            return vector<int>(nums.size());
        }
        if (zero_cnt == 1) {
            vector<int> ret(nums.size());
            ret[zero_idx[0]] = prod;
            return ret;
        }
        vector<int> ret(nums.size());
        ret[0] = dp2[1];
        ret[ret.size() - 1] = dp1[dp1.size() - 2];
        for (size_t i = 1; i < nums.size() - 1; i++) {
            ret[i] = dp1[i - 1] * dp2[i + 1];
        }
        return ret;
    }
};