Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# Wrapper for the shortest augmenting path algorithm for solving the 

2# rectangular linear sum assignment problem. The original code was an 

3# implementation of the Hungarian algorithm (Kuhn-Munkres) taken from 

4# scikit-learn, based on original code by Brian Clapper and adapted to NumPy 

5# by Gael Varoquaux. Further improvements by Ben Root, Vlad Niculae, Lars 

6# Buitinck, and Peter Larsen. 

7# 

8# Copyright (c) 2008 Brian M. Clapper <bmc@clapper.org>, Gael Varoquaux 

9# Author: Brian M. Clapper, Gael Varoquaux 

10# License: 3-clause BSD 

11 

12import numpy as np 

13from . import _lsap_module 

14 

15 

16def linear_sum_assignment(cost_matrix, maximize=False): 

17 """Solve the linear sum assignment problem. 

18 

19 The linear sum assignment problem is also known as minimum weight matching 

20 in bipartite graphs. A problem instance is described by a matrix C, where 

21 each C[i,j] is the cost of matching vertex i of the first partite set 

22 (a "worker") and vertex j of the second set (a "job"). The goal is to find 

23 a complete assignment of workers to jobs of minimal cost. 

24 

25 Formally, let X be a boolean matrix where :math:`X[i,j] = 1` iff row i is 

26 assigned to column j. Then the optimal assignment has cost 

27 

28 .. math:: 

29 \\min \\sum_i \\sum_j C_{i,j} X_{i,j} 

30 

31 where, in the case where the matrix X is square, each row is assigned to 

32 exactly one column, and each column to exactly one row. 

33 

34 This function can also solve a generalization of the classic assignment 

35 problem where the cost matrix is rectangular. If it has more rows than 

36 columns, then not every row needs to be assigned to a column, and vice 

37 versa. 

38 

39 Parameters 

40 ---------- 

41 cost_matrix : array 

42 The cost matrix of the bipartite graph. 

43 

44 maximize : bool (default: False) 

45 Calculates a maximum weight matching if true. 

46 

47 Returns 

48 ------- 

49 row_ind, col_ind : array 

50 An array of row indices and one of corresponding column indices giving 

51 the optimal assignment. The cost of the assignment can be computed 

52 as ``cost_matrix[row_ind, col_ind].sum()``. The row indices will be 

53 sorted; in the case of a square cost matrix they will be equal to 

54 ``numpy.arange(cost_matrix.shape[0])``. 

55 

56 Notes 

57 ----- 

58 .. versionadded:: 0.17.0 

59 

60 References 

61 ---------- 

62 

63 1. https://en.wikipedia.org/wiki/Assignment_problem 

64 

65 2. DF Crouse. On implementing 2D rectangular assignment algorithms. 

66 *IEEE Transactions on Aerospace and Electronic Systems*, 

67 52(4):1679-1696, August 2016, https://doi.org/10.1109/TAES.2016.140952 

68 

69 Examples 

70 -------- 

71 >>> cost = np.array([[4, 1, 3], [2, 0, 5], [3, 2, 2]]) 

72 >>> from scipy.optimize import linear_sum_assignment 

73 >>> row_ind, col_ind = linear_sum_assignment(cost) 

74 >>> col_ind 

75 array([1, 0, 2]) 

76 >>> cost[row_ind, col_ind].sum() 

77 5 

78 """ 

79 cost_matrix = np.asarray(cost_matrix) 

80 if len(cost_matrix.shape) != 2: 

81 raise ValueError("expected a matrix (2-D array), got a %r array" 

82 % (cost_matrix.shape,)) 

83 

84 if not (np.issubdtype(cost_matrix.dtype, np.number) or 

85 cost_matrix.dtype == np.dtype(np.bool)): 

86 raise ValueError("expected a matrix containing numerical entries, got %s" 

87 % (cost_matrix.dtype,)) 

88 

89 if maximize: 

90 cost_matrix = -cost_matrix 

91 

92 if np.any(np.isneginf(cost_matrix) | np.isnan(cost_matrix)): 

93 raise ValueError("matrix contains invalid numeric entries") 

94 

95 cost_matrix = cost_matrix.astype(np.double) 

96 a = np.arange(np.min(cost_matrix.shape)) 

97 

98 # The algorithm expects more columns than rows in the cost matrix. 

99 if cost_matrix.shape[1] < cost_matrix.shape[0]: 

100 b = _lsap_module.calculate_assignment(cost_matrix.T) 

101 indices = np.argsort(b) 

102 return (b[indices], a[indices]) 

103 else: 

104 b = _lsap_module.calculate_assignment(cost_matrix) 

105 return (a, b)