LintCode/C++/post-office-problem.cpp at master · kamyu104/LintCode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
// Time:  O(k * n^2)
// Space: O(n)

class Solution {
public:
    /**
     * @param A an integer array
     * @param k an integer
     * @return an integer
     */
    int postOffice(vector<int>& A, int k) {
        const int n = A.size();
        if (A.empty() || k >= n) {
            return 0;
        }

        sort(A.begin(), A.end());  // Time:  O(nlogn)

        // Precompute cost.
        // Time:  O(n)
        // Space: O(n)
        vector<int> before_cost(n), after_cost(n);
        computeCost(A, &before_cost, &after_cost);

        // DP of sum.
        // Time:  O(k * n^2)
        // Space: O(n)
        // sum[i][j] denotes the smallest sum of
        // picking i post offices for the first j houses.
        vector<vector<int>> sum(2, vector<int>(A.size() + 1, INT_MAX));
        sum[0][0] = 0;
        for (int i = 1; i <= k; ++i) {
            for (int j = 0; j <= n; ++j) {
                // We can skip this line due to sum[i][j] <= sum[i - 2][j]
                // sum[i % 2][j] = INT_MAX;
                for (int r = 1; r <= j; ++r) {
                    if (sum[(i - 1) % 2][j - r] != INT_MAX) {
                        sum[i % 2][j] = min(sum[i % 2][j],
                                            sum[(i - 1) % 2][j - r] +
                                            minCost(A,
                                                    before_cost, after_cost,
                                                    j - r, j - 1));
                    }
                }
            }
        }
        return sum[k % 2][n];
    }

    void computeCost(const vector<int>& A,
                     vector<int> *before_cost, vector<int> *after_cost) {
        const int n = A.size();
        for (int i = 1; i < n; ++i) {
            // before_cost[i] = Sum(A[i] - A[k]) for k < i
            (*before_cost)[i] = (*before_cost)[i - 1] +
                                i * (A[i] - A[i - 1]);
        }
        for (int i = n - 2; i >= 0; --i) {
            // after_cost[i] = Sum(A[k] - A[i]) for k > i
            (*after_cost)[i] = (*after_cost)[i + 1] +
                               (n - 1 - i) * (A[i + 1] - A[i]);
        }
    }

    int minCost(const vector<int>& A,
                const vector<int> &before_cost, const vector<int> &after_cost,
                const int i, const int j) {
            // Min cost of building a post office between house (i, j).
            // This post office must be in median position.
            const int n = A.size();
            int mid = i + (j - i) / 2;
            int before_mid = before_cost[mid] -
                             (before_cost[i] + i * (A[mid] - A[i]));
            int after_mid = after_cost[mid] -
                            (after_cost[j] + (n - 1 - j) * (A[j] - A[mid]));
            return before_mid + after_mid;
    }
};

// Time:  O(n^3)
// Space: O(n^2)
class Solution2 {
public:
    /**
     * @param A an integer array
     * @param k an integer
     * @return an integer
     */
    int postOffice(vector<int>& A, int k) {
        const int n = A.size();
        if (A.empty() || k >= n) {
            return 0;
        }

        sort(A.begin(), A.end());  // Time:  O(nlogn)

        // Precompute cost.
        // Time:  O(n^3)
        // Space: O(n^2)
        vector<vector<int>> cost(A.size() + 1, vector<int>(A.size() + 1, 0));
        computeMinCost(A, &cost);

        // DP of sum.
        // Time:  O(k * n^2)
        // Space: O(k * n)
        // sum[i][j] denotes the smallest sum of
        // picking i post offices for the first j houses.
        vector<vector<int>> sum(k + 1, vector<int>(A.size() + 1, INT_MAX));
        sum[0][0] = 0;
        for (int i = 1; i <= k; ++i) {
            for (int j = 0; j < n; ++j) {
                if (sum[i - 1][j] != INT_MAX) {
                    for (int r = 1; j + r <= n; ++r) {
                       sum[i][j + r] = min(sum[i][j + r],
                                           sum[i - 1][j] + cost[j + 1][j + r]);
                    }
                }
            }
        }
        return sum[k][n];
    }

    void computeMinCost(const vector<int>& A, vector<vector<int>> *cost) {
        // Min cost of building a post office between house (i, j).
        // This post office must be in median position.
        const int n = A.size();
        for (int i = 0; i < n; ++i) {
            for (int j = i; j < n; ++j) {
                int mid = (i + j) / 2;
                for (int r = i; r <= mid; ++r) {
                    (*cost)[i + 1][j + 1] += A[mid] - A[r];
                }
                for (int r = mid + 1; r <= j; ++r) {
                    (*cost)[i + 1][j + 1] += A[r] - A[mid];
                }
            }
        }
    }
};